From 9d5b66af5f88f0e20db291e81b54d7203253b208 Mon Sep 17 00:00:00 2001 From: rockdu Date: Sun, 2 Aug 2026 20:03:17 -0700 Subject: [PATCH 01/50] feat(fsdp): compile per-parameter precision plans onto FSDP2 with autocast forward --- miles/backends/fsdp_utils/actor.py | 97 ++++--- miles/backends/fsdp_utils/configs/ltx.py | 7 +- .../configs/train_pipeline_config.py | 3 + .../backends/fsdp_utils/loss_hub/flow_grpo.py | 17 +- miles/backends/fsdp_utils/loss_hub/nft.py | 7 +- miles/backends/fsdp_utils/loss_hub/utils.py | 8 - miles/backends/fsdp_utils/precision.py | 237 ++++++++++++++++++ miles/utils/arguments.py | 38 ++- .../fsdp_utils/test_precision_plan.py | 100 ++++++++ 9 files changed, 438 insertions(+), 76 deletions(-) delete mode 100644 miles/backends/fsdp_utils/loss_hub/utils.py create mode 100644 miles/backends/fsdp_utils/precision.py create mode 100644 tests/fast/backends/fsdp_utils/test_precision_plan.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d9430378..198c5413 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -39,6 +39,13 @@ from .lr_scheduler import get_lr_scheduler from .metrics import new_metric_buffer from .parallel import create_fsdp_parallel_state +from .precision import ( + clip_grad_norm_mixed, + compile_precision_plan, + log_precision_summary, + resolve_dtype, + sync_replicated_grads, +) from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -94,8 +101,8 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.prof = TrainProfiler(args) - self._master_dtype = _resolve_dtype(args.fsdp_master_dtype) - self._forward_dtype = _resolve_dtype(args.diffusion_forward_dtype) + self._master_dtype = resolve_dtype(args.fsdp_master_dtype) + self._forward_dtype = resolve_dtype(args.diffusion_forward_dtype) from miles.utils.misc import load_function @@ -110,6 +117,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty materialize_weights = rank == 0 self.models: dict[str, torch.nn.Module] = {} + self._noshard_params: list[torch.nn.Parameter] = [] for component in args.update_weight_target_modules: # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap with self._model_init_context(materialize_weights=materialize_weights): @@ -128,6 +136,17 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if args.gradient_checkpointing: self.model_backend.enable_gradient_checkpointing(model) + # Resolve the family precision plan on clean FQNs (pre-LoRA, pre-FSDP). + compiled_precision = compile_precision_plan( + model, + self.train_pipeline_config.precision_spec, + default_dtype=self._forward_dtype, + ) + compiled_precision.apply_master_casts() + if rank == 0: + log_precision_summary(component, compiled_precision, default_dtype=self._forward_dtype) + self._noshard_params.extend(plan.tensor for plan in compiled_precision.noshard_params) + if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) @@ -143,6 +162,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), + ignored_params=compiled_precision.ignored_params(), ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -464,11 +484,16 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: self.prof.step(rollout_id=rollout_id) if not self.args.debug_skip_optimizer_step: self.scaler.unscale_(self.optimizer) - grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) - if isinstance(grad_norm, DTensor): - # clip returns a lazily-reduced partial norm; materialize it, - # otherwise the logged metric leaks the local shard's value. - grad_norm = grad_norm.full_tensor() + if self._noshard_params: + # FSDP ignores no-shard params; average their rank-local grads before clipping. + sync_replicated_grads(self._noshard_params, self.parallel_state.get_mesh("fsdp")) + grad_norm = clip_grad_norm_mixed(self.model.parameters(), self.args.clip_grad) + else: + grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) + if isinstance(grad_norm, DTensor): + # clip returns a lazily-reduced partial norm; materialize it, + # otherwise the logged metric leaks the local shard's value. + grad_norm = grad_norm.full_tensor() metrics.emit_replicated("grad_norm", grad_norm) self.scaler.step(self.optimizer) self.scaler.update() @@ -512,22 +537,21 @@ def _forward_train_pair_batch( train_pipeline_config = self.train_pipeline_config forward_dtype = self._forward_dtype - latents_input = prepared.latents.to(forward_dtype) - timesteps_input = prepared.timesteps_for_model.to(forward_dtype) - def _compute_noise_pred() -> torch.Tensor: - return train_pipeline_config.compute_noise_pred( - model=prepared.model, - latents_input=latents_input, - timesteps_input=timesteps_input, - pos_cond=prepared.pos_cond, - neg_cond=prepared.neg_cond, - joint_cond=prepared.joint_cond, - use_cfg=prepared.use_cfg, - cfg_batching=prepared.cfg_batching, - guidance_scale=prepared.guidance_scale, - true_cfg_scale=prepared.true_cfg_scale, - ) + # Inputs/params keep their resident dtypes; compute dtype is autocast-managed. + with torch.autocast("cuda", dtype=forward_dtype, enabled=forward_dtype != torch.float32): + return train_pipeline_config.compute_noise_pred( + model=prepared.model, + latents_input=prepared.latents, + timesteps_input=prepared.timesteps_for_model, + pos_cond=prepared.pos_cond, + neg_cond=prepared.neg_cond, + joint_cond=prepared.joint_cond, + use_cfg=prepared.use_cfg, + cfg_batching=prepared.cfg_batching, + guidance_scale=prepared.guidance_scale, + true_cfg_scale=prepared.true_cfg_scale, + ) new_pred = _compute_noise_pred() @@ -580,10 +604,6 @@ def move_torch_optimizer(optimizer, device): torch.cuda.synchronize() -def _resolve_dtype(name: str) -> torch.dtype: - return {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}[name] - - def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) -> torch.nn.Module: """Apply PEFT LoRA, leaving non-rank0 adapters uninitialized on meta.""" from peft import LoraConfig, get_peft_model @@ -609,9 +629,11 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None): +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, ignored_params=None): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard + if cpu_offload and ignored_params: + raise ValueError("no-shard precision rules are incompatible with --fsdp-cpu-offload") offload_policy = CPUOffloadPolicy() if cpu_offload else None layer_cls_to_wrap = no_split_modules if no_split_modules is not None else model._no_split_modules @@ -619,29 +641,26 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules modules = [module for name, module in model.named_modules() if module.__class__.__name__ in layer_cls_to_wrap] - param_dtype = _resolve_dtype(args.diffusion_forward_dtype) - reduce_dtype = _resolve_dtype(args.fsdp_reduce_dtype) + param_dtype = resolve_dtype(args.diffusion_forward_dtype) + reduce_dtype = resolve_dtype(args.fsdp_reduce_dtype) logger.info( - f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, reduce_dtype={reduce_dtype}" + f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, " + f"reduce_dtype={reduce_dtype}, ignored_params={len(ignored_params) if ignored_params else 0}" ) fsdp_kwargs = { + # Inputs are not cast: compute dtype comes from the trainer's autocast, which + # also keeps grad-ckpt recompute dtypes consistent with the forward. "mp_policy": MixedPrecisionPolicy( param_dtype=param_dtype, reduce_dtype=reduce_dtype, + cast_forward_inputs=False, ), "offload_policy": offload_policy, "mesh": mesh, } - - if args.gradient_checkpointing: - # MixedPrecisionPolicy does not cast buffers; a buffer above param_dtype - # makes the ckpt recompute dtype-diverge from the forward and abort. - for module in model.modules(): - for name, buf in module.named_buffers(recurse=False): - if buf.is_floating_point() and buf.dtype != param_dtype: - persistent = name not in module._non_persistent_buffers_set - module.register_buffer(name, buf.to(param_dtype), persistent=persistent) + if ignored_params: + fsdp_kwargs["ignored_params"] = ignored_params for module in modules: fully_shard(module, **fsdp_kwargs) diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index a4f77e7a..91271f88 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -129,7 +129,6 @@ def forward_velocity( from ltx_core.model.transformer.modality import Modality from ltx_core.utils import to_denoised - device = latents_input.device dtype = latents_input.dtype B = latents_input.shape[0] @@ -148,10 +147,8 @@ def forward_velocity( context=cond["context"].to(dtype), context_mask=None, ) - # FSDP mixed precision casts parameters but does not replace LTX's - # operation-level autocast semantics. - with torch.autocast(device_type=str(device).split(":")[0], dtype=dtype): - velocity, _ = model(video=video_modality, audio=None, perturbations=None) + # Compute dtype comes from the trainer's ambient autocast around compute_noise_pred. + velocity, _ = model(video=video_modality, audio=None, perturbations=None) # Keep the original fp32 denoised reconstruction path: although this is # algebraically an identity for T2V, strict e2e metrics depend on its rounding. diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index c392ea57..049474e2 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -18,6 +18,7 @@ import torch from miles.utils.types import CondKwargs +from ..precision import PrecisionSpec _REGISTRY: dict[str, type[TrainPipelineConfig]] = {} @@ -82,6 +83,8 @@ class TrainPipelineConfig(abc.ABC): supports_cfg_training: bool = True # Rollout parity patch group applied by the engine (see monkey_patches; None = none). rollout_patch_group: str | None = None + # Weight-precision rules (master/gather dtypes) compiled onto FSDP2; see precision.py. + precision_spec: PrecisionSpec = PrecisionSpec() # Default component paths (miles custom-function style); CLI args override. model_backend_path: str = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" # Native model package import path; required when model_backend_path is MilesModelBackend. diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 00f787b6..162cb7d2 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -5,7 +5,6 @@ import torch from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch -from miles.backends.fsdp_utils.loss_hub.utils import cast_cond_to_dtype from miles.backends.fsdp_utils.metrics import record_rollout_train_abs_diff from miles.utils.metric_buffer import MetricBuffer from miles.utils.train_data_utils import stack_train_pair_rollout_debug @@ -74,23 +73,15 @@ def prepare_flow_grpo_batch( if use_cfg else None ) + # Cond tensors keep their rollout dtypes; the DiT forward runs under autocast. cfg_batching = use_cfg and bool(args.fsdp_cfg_batching) joint_cond = pos_cond = neg_cond = None if cfg_batching: - joint_cond = cast_cond_to_dtype( - config.collate_cond_for_sample_batch(pos_list + neg_list, device, pad_to_len=pad_to_len), - ctx.forward_dtype, - ) + joint_cond = config.collate_cond_for_sample_batch(pos_list + neg_list, device, pad_to_len=pad_to_len) else: - pos_cond = cast_cond_to_dtype( - config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len), - ctx.forward_dtype, - ) + pos_cond = config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len) if use_cfg and neg_list is not None: - neg_cond = cast_cond_to_dtype( - config.collate_cond_for_sample_batch(neg_list, device, pad_to_len=pad_to_len), - ctx.forward_dtype, - ) + neg_cond = config.collate_cond_for_sample_batch(neg_list, device, pad_to_len=pad_to_len) return PreparedBatch( latents=latents, diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 183bd47f..51e2c678 100644 --- a/miles/backends/fsdp_utils/loss_hub/nft.py +++ b/miles/backends/fsdp_utils/loss_hub/nft.py @@ -5,7 +5,6 @@ import torch from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch -from miles.backends.fsdp_utils.loss_hub.utils import cast_cond_to_dtype from miles.utils.metric_buffer import MetricBuffer @@ -38,10 +37,8 @@ def prepare_nft_batch( component_name, model = next(iter(ctx.models.items())) pos_list = [config.prepare_cond_kwargs(batch[i]["denoising_env"].pos_cond_kwargs, device) for i in range(bsz)] - pos_cond = cast_cond_to_dtype( - config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len), - ctx.forward_dtype, - ) + # Cond tensors keep their rollout dtypes; the DiT forward runs under autocast. + pos_cond = config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len) num_train_timesteps = ctx.scheduler.config.num_train_timesteps if config.needs_timestep_scaling: diff --git a/miles/backends/fsdp_utils/loss_hub/utils.py b/miles/backends/fsdp_utils/loss_hub/utils.py deleted file mode 100644 index 3f1d0893..00000000 --- a/miles/backends/fsdp_utils/loss_hub/utils.py +++ /dev/null @@ -1,8 +0,0 @@ -import torch - - -def cast_cond_to_dtype(cond: dict, dtype: torch.dtype) -> dict: - return { - key: value.to(dtype=dtype) if isinstance(value, torch.Tensor) and value.dtype.is_floating_point else value - for key, value in cond.items() - } diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py new file mode 100644 index 00000000..56f2bac7 --- /dev/null +++ b/miles/backends/fsdp_utils/precision.py @@ -0,0 +1,237 @@ +"""Fine-grained weight-precision control for FSDP2. + +A family declares per-tensor dtype intent as PrecisionSpec rules on its +TrainPipelineConfig (last matching rule wins per axis): + - master: resident dtype of the param/buffer (optimizer state precision) + - gather: dtype the param is cast to for FSDP all-gather + forward +Compute dtype is deliberately not managed here: the trainer runs the DiT +forward under torch.autocast(default dtype); op-level exceptions belong to +the monkey-patch registry. + +``compile_precision_plan`` resolves rules to a per-tensor plan and lowers it +onto what FSDP2 can express: + - inline (default): tensor follows its wrap unit's MixedPrecisionPolicy + - no_shard: gather pinned away from the default -> the param is excluded + from sharding (fully_shard ignored_params), replicated at master dtype; + its rank-local grads are averaged manually each step + - gather != master needs sub-shard lowering (a nested fully_shard with its + own policy) and is rejected until a use case exists +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from fnmatch import fnmatch + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor + +logger = logging.getLogger(__name__) + +_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16} + + +def resolve_dtype(name: str) -> torch.dtype: + return _DTYPES[name] + + +# --------------------------------------------------------------------------- +# Spec: per-family declaration (see TrainPipelineConfig.precision_spec) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ModuleSel: + """Matches every param/buffer under modules with fnmatch-ing FQN or exact class name.""" + + fqn: str | None = None + cls: str | None = None + + +@dataclass(frozen=True) +class ParamSel: + """Matches individual params/buffers by fnmatch on their FQN.""" + + fqn: str + + +@dataclass(frozen=True) +class Rule: + """Axes take a dtype name ("fp32"/"bf16"/"fp16"), "default" (the run's default dtype), or None (untouched).""" + + select: ModuleSel | ParamSel + master: str | None = None + gather: str | None = None + + +@dataclass(frozen=True) +class PrecisionSpec: + rules: tuple[Rule, ...] = () + + +# --------------------------------------------------------------------------- +# Compiler: per-tensor plan -> FSDP2 lowering (master casts + no-shard set) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class TensorPlan: + fqn: str + tensor: torch.Tensor + is_buffer: bool + master: torch.dtype | None + gather: torch.dtype | None + + +@dataclass +class CompiledPrecision: + master_casts: list[TensorPlan] + noshard_params: list[TensorPlan] + + def ignored_params(self) -> set[torch.nn.Parameter] | None: + return {plan.tensor for plan in self.noshard_params} or None + + def apply_master_casts(self) -> None: + for plan in self.master_casts: + plan.tensor.data = plan.tensor.data.to(plan.master) + + +def _validate_rule(rule: Rule) -> None: + if rule.master is None and rule.gather is None: + raise ValueError(f"precision rule sets no dtype axis: {rule}") + if isinstance(rule.select, ModuleSel) and rule.select.fqn is None and rule.select.cls is None: + raise ValueError(f"precision rule has an empty ModuleSel: {rule}") + for axis in (rule.master, rule.gather): + if axis is not None and axis != "default" and axis not in _DTYPES: + raise ValueError(f"precision rule has unknown dtype {axis!r}: {rule}") + + +def _matched_module_prefixes(sel: ModuleSel, model: torch.nn.Module) -> list[str]: + prefixes = [] + for mod_fqn, module in model.named_modules(): + if sel.cls is not None and type(module).__name__ != sel.cls: + continue + if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): + continue + prefixes.append(mod_fqn) + return prefixes + + +def compile_precision_plan( + model: torch.nn.Module, + spec: PrecisionSpec, + *, + default_dtype: torch.dtype, +) -> CompiledPrecision: + """Resolve spec rules against the (pre-LoRA, pre-FSDP) model and lower them. + + Raises on rules that match nothing (likely a typo'd pattern) and on plans + FSDP2 cannot express (see module docstring). + """ + for rule in spec.rules: + _validate_rule(rule) + + named: list[tuple[str, torch.Tensor, bool]] = [] + for mod_fqn, module in model.named_modules(): + prefix = f"{mod_fqn}." if mod_fqn else "" + for name, param in module.named_parameters(recurse=False): + named.append((f"{prefix}{name}", param, False)) + for name, buf in module.named_buffers(recurse=False): + named.append((f"{prefix}{name}", buf, True)) + + matchers = [] + for rule in spec.rules: + if isinstance(rule.select, ParamSel): + pattern = rule.select.fqn + matchers.append(lambda fqn, pattern=pattern: fnmatch(fqn, pattern)) + else: + prefixes = _matched_module_prefixes(rule.select, model) + matchers.append(lambda fqn, prefixes=prefixes: any(p == "" or fqn.startswith(f"{p}.") for p in prefixes)) + + def _resolve_axis(axis: str | None) -> torch.dtype | None: + if axis is None: + return None + return default_dtype if axis == "default" else _DTYPES[axis] + + match_counts = [0] * len(spec.rules) + master_casts: list[TensorPlan] = [] + noshard_params: list[TensorPlan] = [] + for fqn, tensor, is_buffer in named: + master = gather = None + for i, rule in enumerate(spec.rules): + if not matchers[i](fqn): + continue + match_counts[i] += 1 + master = rule.master if rule.master is not None else master + gather = rule.gather if rule.gather is not None else gather + if master is None and gather is None: + continue + + plan = TensorPlan(fqn, tensor, is_buffer, _resolve_axis(master), _resolve_axis(gather)) + if plan.gather is not None: + if is_buffer: + raise ValueError(f"precision rule pins gather dtype on buffer {fqn}; buffers are never gathered") + if plan.gather != default_dtype: + effective_master = plan.master if plan.master is not None else tensor.dtype + if effective_master != plan.gather: + raise ValueError( + f"{fqn}: gather={plan.gather} != master={effective_master}; sub-shard lowering " + f"is not implemented, pin master to the same dtype to use no-shard" + ) + noshard_params.append(plan) + if plan.master is not None and plan.master != tensor.dtype: + master_casts.append(plan) + + for rule, count in zip(spec.rules, match_counts, strict=True): + if count == 0: + raise ValueError(f"precision rule matched no tensor: {rule}") + + return CompiledPrecision(master_casts=master_casts, noshard_params=noshard_params) + + +def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: + noshard_bytes = sum(plan.tensor.numel() * plan.tensor.element_size() for plan in compiled.noshard_params) + logger.info( + f"precision[{component}]: default gather dtype {default_dtype}, " + f"{len(compiled.master_casts)} master casts, " + f"{len(compiled.noshard_params)} no-shard params ({noshard_bytes / 1e6:.1f} MB replicated/rank)" + ) + for plan in compiled.master_casts: + logger.info(f"precision[{component}]: master {plan.fqn} -> {plan.master}") + for plan in compiled.noshard_params: + logger.info(f"precision[{component}]: no-shard {plan.fqn} @ {plan.gather}") + + +# --------------------------------------------------------------------------- +# Runtime helpers for no-shard (FSDP-ignored, replicated) params +# --------------------------------------------------------------------------- + + +def sync_replicated_grads(params: list[torch.nn.Parameter], mesh: DeviceMesh) -> None: + """Average rank-local grads of FSDP-ignored params over every dim FSDP reduces over.""" + grads = [p.grad for p in params if p.grad is not None] + for dim in range(mesh.ndim): + group = mesh.get_group(dim) + for grad in grads: + dist.all_reduce(grad, op=dist.ReduceOp.AVG, group=group) + + +def clip_grad_norm_mixed(parameters, max_norm: float) -> torch.Tensor: + """Global grad-norm clip over mixed DTensor + plain grads (torch rejects the mix in one call).""" + from torch.nn.utils import clip_grads_with_norm_, get_total_norm + + params = [p for p in parameters if p.grad is not None] + sharded = [p for p in params if isinstance(p.grad, DTensor)] + replicated = [p for p in params if not isinstance(p.grad, DTensor)] + norms = [] + if sharded: + norms.append(get_total_norm([p.grad for p in sharded]).full_tensor().float()) + if replicated: + norms.append(get_total_norm([p.grad for p in replicated]).float()) + total_norm = torch.linalg.vector_norm(torch.stack(norms)) + clip_grads_with_norm_(sharded, max_norm, total_norm) + clip_grads_with_norm_(replicated, max_norm, total_norm) + return total_norm diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index e0c9c4ac..4e78b822 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -189,17 +189,28 @@ def add_train_arguments(parser): "of bf16 add-non-associativity noise across ranks." ), ) + parser.add_argument( + "--precision-default-dtype", + type=str, + default=None, + choices=["fp16", "bf16", "fp32"], + help=( + "Default dtype for every module the family PrecisionSpec does not pin " + "explicitly: sets both the training-side forward/gather dtype " + "(--diffusion-forward-dtype) and the rollout engine's " + "--sglang-dit-precision. Specific flags win over this default." + ), + ) parser.add_argument( "--diffusion-forward-dtype", type=str, - default="bf16", + default=None, choices=["fp16", "bf16", "fp32"], help=( - "dtype for the DiT forward compute. Used in three places " - "with the same value: sglang-d rollout engine, FSDP " - "MixedPrecisionPolicy.param_dtype on the training side, " - "and the training-side input cast that matches rollout " - "for log-prob alignment." + "dtype for the DiT forward on the training side: FSDP " + "MixedPrecisionPolicy.param_dtype and the torch.autocast the " + "trainer wraps the forward in. Defaults to " + "--precision-default-dtype, else bf16." ), ) parser.add_argument( @@ -1519,6 +1530,21 @@ def miles_validate_args(args): if args.eval_reward_key is None: args.eval_reward_key = args.reward_key + # Resolve the precision default-dtype cascade: specific flag > --precision-default-dtype > built-in. + if args.diffusion_forward_dtype is None: + args.diffusion_forward_dtype = args.precision_default_dtype or "bf16" + if args.precision_default_dtype is not None: + from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig + + sglang_dit = getattr(args, "sglang_dit_precision", None) + if sglang_dit is None or sglang_dit == getattr(PipelineConfig, "dit_precision", None): + args.sglang_dit_precision = args.precision_default_dtype + elif sglang_dit != args.precision_default_dtype: + raise ValueError( + f"--sglang-dit-precision {sglang_dit} conflicts with " + f"--precision-default-dtype {args.precision_default_dtype}" + ) + args.update_weight_target_modules = [ name.strip() for name in args.update_weight_target_module.split(",") if name.strip() ] diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py new file mode 100644 index 00000000..66f8d32d --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -0,0 +1,100 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) + +import pytest +import torch +import torch.nn as nn + +from miles.backends.fsdp_utils.precision import ( + ModuleSel, + ParamSel, + PrecisionSpec, + Rule, + compile_precision_plan, +) + + +class Block(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + self.norm = nn.LayerNorm(8) + self.register_buffer("freqs", torch.zeros(4)) + + +class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([Block(), Block()]) + + +def _model(dtype=torch.bfloat16): + return Tiny().to(dtype) + + +def test_empty_spec_compiles_to_nothing(): + compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) + assert compiled.master_casts == [] + assert compiled.noshard_params == [] + assert compiled.ignored_params() is None + + +def test_module_cls_rule_pins_norms_to_noshard(): + model = _model() + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"),)) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert {p.fqn for p in compiled.noshard_params} == { + f"blocks.{i}.norm.{n}" for i in range(2) for n in ("weight", "bias") + } + assert all(p.master is torch.float32 for p in compiled.master_casts) + compiled.apply_master_casts() + assert model.blocks[0].norm.weight.dtype is torch.float32 + assert model.blocks[0].linear.weight.dtype is torch.bfloat16 + assert compiled.ignored_params() == {p for b in model.blocks for p in b.norm.parameters()} + + +def test_param_fqn_rule_casts_buffer_master(): + model = _model() + spec = PrecisionSpec(rules=(Rule(ParamSel("*.freqs"), master="fp32"),)) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert {p.fqn for p in compiled.master_casts} == {"blocks.0.freqs", "blocks.1.freqs"} + assert compiled.noshard_params == [] + compiled.apply_master_casts() + assert model.blocks[0].freqs.dtype is torch.float32 + + +def test_gather_matching_default_is_inline(): + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert compiled.noshard_params == [] + assert compiled.master_casts == [] + + +def test_gather_diverging_from_master_rejected(): + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) + with pytest.raises(ValueError, match="sub-shard"): + compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + + +def test_gather_on_buffer_rejected(): + spec = PrecisionSpec(rules=(Rule(ParamSel("*.freqs"), master="fp32", gather="fp32"),)) + with pytest.raises(ValueError, match="buffer"): + compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + + +def test_last_rule_wins_per_axis(): + spec = PrecisionSpec( + rules=( + Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"), + Rule(ParamSel("blocks.0.norm.*"), master="default", gather="default"), + ) + ) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert {p.fqn for p in compiled.noshard_params} == {f"blocks.1.norm.{n}" for n in ("weight", "bias")} + + +def test_unmatched_rule_rejected(): + spec = PrecisionSpec(rules=(Rule(ParamSel("no.such.param"), master="fp32"),)) + with pytest.raises(ValueError, match="matched no tensor"): + compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) From 5af4bfeb13ec8a3ef2870321ed2cf6c766c99ce4 Mon Sep 17 00:00:00 2001 From: rockdu Date: Sun, 2 Aug 2026 20:53:52 -0700 Subject: [PATCH 02/50] refactor(fsdp): lower gather overrides to grouped nested fully_shard instead of ignored_params --- miles/backends/fsdp_utils/actor.py | 67 +++++----- miles/backends/fsdp_utils/precision.py | 121 +++++++++--------- .../fsdp_utils/test_precision_plan.py | 49 ++++--- 3 files changed, 121 insertions(+), 116 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 198c5413..e1cd4c38 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -39,13 +39,7 @@ from .lr_scheduler import get_lr_scheduler from .metrics import new_metric_buffer from .parallel import create_fsdp_parallel_state -from .precision import ( - clip_grad_norm_mixed, - compile_precision_plan, - log_precision_summary, - resolve_dtype, - sync_replicated_grads, -) +from .precision import compile_precision_plan, log_precision_summary, resolve_dtype from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -117,7 +111,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty materialize_weights = rank == 0 self.models: dict[str, torch.nn.Module] = {} - self._noshard_params: list[torch.nn.Parameter] = [] for component in args.update_weight_target_modules: # per raw component (wan2.2 has two transformers), before LoRA/FSDP wrap with self._model_init_context(materialize_weights=materialize_weights): @@ -145,7 +138,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty compiled_precision.apply_master_casts() if rank == 0: log_precision_summary(component, compiled_precision, default_dtype=self._forward_dtype) - self._noshard_params.extend(plan.tensor for plan in compiled_precision.noshard_params) if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) @@ -162,7 +154,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), - ignored_params=compiled_precision.ignored_params(), + subshard_groups=compiled_precision.subshard_groups, ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -484,16 +476,11 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: self.prof.step(rollout_id=rollout_id) if not self.args.debug_skip_optimizer_step: self.scaler.unscale_(self.optimizer) - if self._noshard_params: - # FSDP ignores no-shard params; average their rank-local grads before clipping. - sync_replicated_grads(self._noshard_params, self.parallel_state.get_mesh("fsdp")) - grad_norm = clip_grad_norm_mixed(self.model.parameters(), self.args.clip_grad) - else: - grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) - if isinstance(grad_norm, DTensor): - # clip returns a lazily-reduced partial norm; materialize it, - # otherwise the logged metric leaks the local shard's value. - grad_norm = grad_norm.full_tensor() + grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) + if isinstance(grad_norm, DTensor): + # clip returns a lazily-reduced partial norm; materialize it, + # otherwise the logged metric leaks the local shard's value. + grad_norm = grad_norm.full_tensor() metrics.emit_replicated("grad_norm", grad_norm) self.scaler.step(self.optimizer) self.scaler.update() @@ -629,11 +616,9 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, ignored_params=None): +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, subshard_groups=None): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard - if cpu_offload and ignored_params: - raise ValueError("no-shard precision rules are incompatible with --fsdp-cpu-offload") offload_policy = CPUOffloadPolicy() if cpu_offload else None layer_cls_to_wrap = no_split_modules if no_split_modules is not None else model._no_split_modules @@ -645,24 +630,32 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules reduce_dtype = resolve_dtype(args.fsdp_reduce_dtype) logger.info( f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, " - f"reduce_dtype={reduce_dtype}, ignored_params={len(ignored_params) if ignored_params else 0}" + f"reduce_dtype={reduce_dtype}, sub-shard groups={len(subshard_groups) if subshard_groups else 0}" ) - fsdp_kwargs = { - # Inputs are not cast: compute dtype comes from the trainer's autocast, which - # also keeps grad-ckpt recompute dtypes consistent with the forward. - "mp_policy": MixedPrecisionPolicy( - param_dtype=param_dtype, - reduce_dtype=reduce_dtype, - cast_forward_inputs=False, - ), - "offload_policy": offload_policy, - "mesh": mesh, - } - if ignored_params: - fsdp_kwargs["ignored_params"] = ignored_params + def _fsdp_kwargs(policy_param_dtype): + return { + # Inputs are not cast: compute dtype comes from the trainer's autocast, which + # also keeps grad-ckpt recompute dtypes consistent with the forward. + "mp_policy": MixedPrecisionPolicy( + param_dtype=policy_param_dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=False, + ), + "offload_policy": offload_policy, + "mesh": mesh, + } + + # Sub-shard groups wrap first so the block/root units exclude their params. + subshard_modules = set() + for group in subshard_groups or (): + fully_shard(group.modules, **_fsdp_kwargs(group.param_dtype)) + subshard_modules.update(group.modules) + fsdp_kwargs = _fsdp_kwargs(param_dtype) for module in modules: + if module in subshard_modules: + raise ValueError(f"{type(module).__name__} is both an FSDP block unit and a precision sub-shard target") fully_shard(module, **fsdp_kwargs) fully_shard(model, **fsdp_kwargs) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 56f2bac7..677627da 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -11,23 +11,22 @@ ``compile_precision_plan`` resolves rules to a per-tensor plan and lowers it onto what FSDP2 can express: - inline (default): tensor follows its wrap unit's MixedPrecisionPolicy - - no_shard: gather pinned away from the default -> the param is excluded - from sharding (fully_shard ignored_params), replicated at master dtype; - its rank-local grads are averaged manually each step - - gather != master needs sub-shard lowering (a nested fully_shard with its - own policy) and is rejected until a use case exists + - sub_shard: gather pinned away from the default -> the owning modules are + grouped into one nested fully_shard unit per gather dtype with its own + MixedPrecisionPolicy, staying fully inside FSDP (DTensor params, FSDP + grad reduction, DCP/offload as usual) at the cost of one extra + all-gather per group per forward +Gather intent that does not align to a module boundary (e.g. a bare +nn.Parameter on a block) cannot be sub-wrapped and is rejected. """ from __future__ import annotations import logging -from dataclasses import dataclass +from dataclasses import dataclass, field from fnmatch import fnmatch import torch -import torch.distributed as dist -from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.tensor import DTensor logger = logging.getLogger(__name__) @@ -73,7 +72,7 @@ class PrecisionSpec: # --------------------------------------------------------------------------- -# Compiler: per-tensor plan -> FSDP2 lowering (master casts + no-shard set) +# Compiler: per-tensor plan -> FSDP2 lowering (master casts + sub-shard groups) # --------------------------------------------------------------------------- @@ -86,13 +85,19 @@ class TensorPlan: gather: torch.dtype | None +@dataclass +class SubShardGroup: + """Modules to wrap as one nested fully_shard unit with param_dtype=gather.""" + + param_dtype: torch.dtype + modules: list[torch.nn.Module] = field(default_factory=list) + param_fqns: list[str] = field(default_factory=list) + + @dataclass class CompiledPrecision: master_casts: list[TensorPlan] - noshard_params: list[TensorPlan] - - def ignored_params(self) -> set[torch.nn.Parameter] | None: - return {plan.tensor for plan in self.noshard_params} or None + subshard_groups: list[SubShardGroup] def apply_master_casts(self) -> None: for plan in self.master_casts: @@ -134,13 +139,16 @@ def compile_precision_plan( for rule in spec.rules: _validate_rule(rule) - named: list[tuple[str, torch.Tensor, bool]] = [] + named: list[tuple[str, torch.Tensor, bool, str, torch.nn.Module]] = [] + float_params_of_module: dict[str, set[str]] = {} for mod_fqn, module in model.named_modules(): prefix = f"{mod_fqn}." if mod_fqn else "" for name, param in module.named_parameters(recurse=False): - named.append((f"{prefix}{name}", param, False)) + named.append((f"{prefix}{name}", param, False, mod_fqn, module)) + if param.is_floating_point(): + float_params_of_module.setdefault(mod_fqn, set()).add(f"{prefix}{name}") for name, buf in module.named_buffers(recurse=False): - named.append((f"{prefix}{name}", buf, True)) + named.append((f"{prefix}{name}", buf, True, mod_fqn, module)) matchers = [] for rule in spec.rules: @@ -158,8 +166,8 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: match_counts = [0] * len(spec.rules) master_casts: list[TensorPlan] = [] - noshard_params: list[TensorPlan] = [] - for fqn, tensor, is_buffer in named: + gather_by_module: dict[str, tuple[torch.nn.Module, dict[str, torch.dtype]]] = {} + for fqn, tensor, is_buffer, mod_fqn, module in named: master = gather = None for i, rule in enumerate(spec.rules): if not matchers[i](fqn): @@ -175,13 +183,9 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: if is_buffer: raise ValueError(f"precision rule pins gather dtype on buffer {fqn}; buffers are never gathered") if plan.gather != default_dtype: - effective_master = plan.master if plan.master is not None else tensor.dtype - if effective_master != plan.gather: - raise ValueError( - f"{fqn}: gather={plan.gather} != master={effective_master}; sub-shard lowering " - f"is not implemented, pin master to the same dtype to use no-shard" - ) - noshard_params.append(plan) + if mod_fqn == "": + raise ValueError(f"{fqn}: cannot sub-wrap the root module for a gather override") + gather_by_module.setdefault(mod_fqn, (module, {}))[1][fqn] = plan.gather if plan.master is not None and plan.master != tensor.dtype: master_casts.append(plan) @@ -189,49 +193,38 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: if count == 0: raise ValueError(f"precision rule matched no tensor: {rule}") - return CompiledPrecision(master_casts=master_casts, noshard_params=noshard_params) + # Gather overrides lower to whole-module nested fully_shard units, so every + # float param of an affected module must agree on one gather dtype. + groups: dict[torch.dtype, SubShardGroup] = {} + for mod_fqn, (module, requests) in gather_by_module.items(): + dtypes = set(requests.values()) + if len(dtypes) > 1: + raise ValueError( + f"module {mod_fqn} mixes gather dtypes {sorted(map(str, dtypes))}; FSDP sub-wrap needs one" + ) + missing = float_params_of_module.get(mod_fqn, set()) - requests.keys() + if missing: + raise ValueError( + f"gather override must cover the whole module {mod_fqn} for FSDP sub-wrap; " + f"missing {sorted(missing)}" + ) + dtype = dtypes.pop() + group = groups.setdefault(dtype, SubShardGroup(param_dtype=dtype)) + group.modules.append(module) + group.param_fqns.extend(sorted(requests)) + + return CompiledPrecision(master_casts=master_casts, subshard_groups=list(groups.values())) def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: - noshard_bytes = sum(plan.tensor.numel() * plan.tensor.element_size() for plan in compiled.noshard_params) logger.info( f"precision[{component}]: default gather dtype {default_dtype}, " - f"{len(compiled.master_casts)} master casts, " - f"{len(compiled.noshard_params)} no-shard params ({noshard_bytes / 1e6:.1f} MB replicated/rank)" + f"{len(compiled.master_casts)} master casts, {len(compiled.subshard_groups)} sub-shard groups" ) for plan in compiled.master_casts: logger.info(f"precision[{component}]: master {plan.fqn} -> {plan.master}") - for plan in compiled.noshard_params: - logger.info(f"precision[{component}]: no-shard {plan.fqn} @ {plan.gather}") - - -# --------------------------------------------------------------------------- -# Runtime helpers for no-shard (FSDP-ignored, replicated) params -# --------------------------------------------------------------------------- - - -def sync_replicated_grads(params: list[torch.nn.Parameter], mesh: DeviceMesh) -> None: - """Average rank-local grads of FSDP-ignored params over every dim FSDP reduces over.""" - grads = [p.grad for p in params if p.grad is not None] - for dim in range(mesh.ndim): - group = mesh.get_group(dim) - for grad in grads: - dist.all_reduce(grad, op=dist.ReduceOp.AVG, group=group) - - -def clip_grad_norm_mixed(parameters, max_norm: float) -> torch.Tensor: - """Global grad-norm clip over mixed DTensor + plain grads (torch rejects the mix in one call).""" - from torch.nn.utils import clip_grads_with_norm_, get_total_norm - - params = [p for p in parameters if p.grad is not None] - sharded = [p for p in params if isinstance(p.grad, DTensor)] - replicated = [p for p in params if not isinstance(p.grad, DTensor)] - norms = [] - if sharded: - norms.append(get_total_norm([p.grad for p in sharded]).full_tensor().float()) - if replicated: - norms.append(get_total_norm([p.grad for p in replicated]).float()) - total_norm = torch.linalg.vector_norm(torch.stack(norms)) - clip_grads_with_norm_(sharded, max_norm, total_norm) - clip_grads_with_norm_(replicated, max_norm, total_norm) - return total_norm + for group in compiled.subshard_groups: + logger.info( + f"precision[{component}]: sub-shard @ {group.param_dtype}: " + f"{len(group.modules)} modules, params {group.param_fqns}" + ) diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 66f8d32d..555eed56 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -36,22 +36,29 @@ def _model(dtype=torch.bfloat16): def test_empty_spec_compiles_to_nothing(): compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.master_casts == [] - assert compiled.noshard_params == [] - assert compiled.ignored_params() is None + assert compiled.subshard_groups == [] -def test_module_cls_rule_pins_norms_to_noshard(): +def test_module_cls_rule_lowers_norms_to_one_subshard_group(): model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert {p.fqn for p in compiled.noshard_params} == { - f"blocks.{i}.norm.{n}" for i in range(2) for n in ("weight", "bias") - } - assert all(p.master is torch.float32 for p in compiled.master_casts) + assert len(compiled.subshard_groups) == 1 + group = compiled.subshard_groups[0] + assert group.param_dtype is torch.float32 + assert group.modules == [model.blocks[0].norm, model.blocks[1].norm] + assert set(group.param_fqns) == {f"blocks.{i}.norm.{n}" for i in range(2) for n in ("weight", "bias")} compiled.apply_master_casts() assert model.blocks[0].norm.weight.dtype is torch.float32 assert model.blocks[0].linear.weight.dtype is torch.bfloat16 - assert compiled.ignored_params() == {p for b in model.blocks for p in b.norm.parameters()} + + +def test_gather_without_master_is_allowed(): + # Nested policy casts at all-gather, so gather may diverge from master. + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert compiled.master_casts == [] + assert compiled.subshard_groups[0].param_dtype is torch.float32 def test_param_fqn_rule_casts_buffer_master(): @@ -59,7 +66,7 @@ def test_param_fqn_rule_casts_buffer_master(): spec = PrecisionSpec(rules=(Rule(ParamSel("*.freqs"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert {p.fqn for p in compiled.master_casts} == {"blocks.0.freqs", "blocks.1.freqs"} - assert compiled.noshard_params == [] + assert compiled.subshard_groups == [] compiled.apply_master_casts() assert model.blocks[0].freqs.dtype is torch.float32 @@ -67,13 +74,24 @@ def test_param_fqn_rule_casts_buffer_master(): def test_gather_matching_default_is_inline(): spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) - assert compiled.noshard_params == [] + assert compiled.subshard_groups == [] assert compiled.master_casts == [] -def test_gather_diverging_from_master_rejected(): - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) - with pytest.raises(ValueError, match="sub-shard"): +def test_partial_module_gather_rejected(): + spec = PrecisionSpec(rules=(Rule(ParamSel("*.norm.weight"), gather="fp32"),)) + with pytest.raises(ValueError, match="whole module"): + compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + + +def test_mixed_gather_dtypes_in_module_rejected(): + spec = PrecisionSpec( + rules=( + Rule(ParamSel("*.norm.weight"), gather="fp32"), + Rule(ParamSel("*.norm.bias"), gather="fp16"), + ) + ) + with pytest.raises(ValueError, match="mixes gather dtypes"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) @@ -84,14 +102,15 @@ def test_gather_on_buffer_rejected(): def test_last_rule_wins_per_axis(): + model = _model() spec = PrecisionSpec( rules=( Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"), Rule(ParamSel("blocks.0.norm.*"), master="default", gather="default"), ) ) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) - assert {p.fqn for p in compiled.noshard_params} == {f"blocks.1.norm.{n}" for n in ("weight", "bias")} + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert compiled.subshard_groups[0].modules == [model.blocks[1].norm] def test_unmatched_rule_rejected(): From 987839f689f7b8850cfee455e77d54e7612ddca1 Mon Sep 17 00:00:00 2001 From: rockdu Date: Sun, 2 Aug 2026 21:04:13 -0700 Subject: [PATCH 03/50] perf(fsdp): keep precision sub-shard groups unsharded through backward --- miles/backends/fsdp_utils/actor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index e1cd4c38..e64408fd 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -647,9 +647,10 @@ def _fsdp_kwargs(policy_param_dtype): } # Sub-shard groups wrap first so the block/root units exclude their params. + # Groups are tiny by design: ZeRO-2 style (no backward re-gather) costs a few MB. subshard_modules = set() for group in subshard_groups or (): - fully_shard(group.modules, **_fsdp_kwargs(group.param_dtype)) + fully_shard(group.modules, reshard_after_forward=False, **_fsdp_kwargs(group.param_dtype)) subshard_modules.update(group.modules) fsdp_kwargs = _fsdp_kwargs(param_dtype) From d6eb22fafb3a625d780a78d3becd9e68c7f74cac Mon Sep 17 00:00:00 2001 From: rockdu Date: Sun, 2 Aug 2026 21:16:22 -0700 Subject: [PATCH 04/50] refactor(fsdp): move precision dtype defaults to set_default_diffusion_args and slim compiler output types --- miles/backends/fsdp_utils/precision.py | 35 +++++++++---------- miles/utils/arguments.py | 30 ++++++++-------- .../fsdp_utils/test_precision_plan.py | 3 +- 3 files changed, 34 insertions(+), 34 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 677627da..a8b5c37e 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -11,7 +11,7 @@ ``compile_precision_plan`` resolves rules to a per-tensor plan and lowers it onto what FSDP2 can express: - inline (default): tensor follows its wrap unit's MixedPrecisionPolicy - - sub_shard: gather pinned away from the default -> the owning modules are + - sub-shard: gather pinned away from the default -> the owning modules are grouped into one nested fully_shard unit per gather dtype with its own MixedPrecisionPolicy, staying fully inside FSDP (DTensor params, FSDP grad reduction, DCP/offload as usual) at the cost of one extra @@ -77,12 +77,10 @@ class PrecisionSpec: @dataclass(frozen=True) -class TensorPlan: +class MasterCast: fqn: str tensor: torch.Tensor - is_buffer: bool - master: torch.dtype | None - gather: torch.dtype | None + dtype: torch.dtype @dataclass @@ -96,12 +94,12 @@ class SubShardGroup: @dataclass class CompiledPrecision: - master_casts: list[TensorPlan] + master_casts: list[MasterCast] subshard_groups: list[SubShardGroup] def apply_master_casts(self) -> None: - for plan in self.master_casts: - plan.tensor.data = plan.tensor.data.to(plan.master) + for cast in self.master_casts: + cast.tensor.data = cast.tensor.data.to(cast.dtype) def _validate_rule(rule: Rule) -> None: @@ -165,7 +163,7 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: return default_dtype if axis == "default" else _DTYPES[axis] match_counts = [0] * len(spec.rules) - master_casts: list[TensorPlan] = [] + master_casts: list[MasterCast] = [] gather_by_module: dict[str, tuple[torch.nn.Module, dict[str, torch.dtype]]] = {} for fqn, tensor, is_buffer, mod_fqn, module in named: master = gather = None @@ -175,19 +173,18 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: match_counts[i] += 1 master = rule.master if rule.master is not None else master gather = rule.gather if rule.gather is not None else gather - if master is None and gather is None: - continue - plan = TensorPlan(fqn, tensor, is_buffer, _resolve_axis(master), _resolve_axis(gather)) - if plan.gather is not None: + gather_dtype = _resolve_axis(gather) + if gather_dtype is not None: if is_buffer: raise ValueError(f"precision rule pins gather dtype on buffer {fqn}; buffers are never gathered") - if plan.gather != default_dtype: + if gather_dtype != default_dtype: if mod_fqn == "": raise ValueError(f"{fqn}: cannot sub-wrap the root module for a gather override") - gather_by_module.setdefault(mod_fqn, (module, {}))[1][fqn] = plan.gather - if plan.master is not None and plan.master != tensor.dtype: - master_casts.append(plan) + gather_by_module.setdefault(mod_fqn, (module, {}))[1][fqn] = gather_dtype + master_dtype = _resolve_axis(master) + if master_dtype is not None and master_dtype != tensor.dtype: + master_casts.append(MasterCast(fqn, tensor, master_dtype)) for rule, count in zip(spec.rules, match_counts, strict=True): if count == 0: @@ -221,8 +218,8 @@ def log_precision_summary(component: str, compiled: CompiledPrecision, *, defaul f"precision[{component}]: default gather dtype {default_dtype}, " f"{len(compiled.master_casts)} master casts, {len(compiled.subshard_groups)} sub-shard groups" ) - for plan in compiled.master_casts: - logger.info(f"precision[{component}]: master {plan.fqn} -> {plan.master}") + for cast in compiled.master_casts: + logger.info(f"precision[{component}]: master {cast.fqn} -> {cast.dtype}") for group in compiled.subshard_groups: logger.info( f"precision[{component}]: sub-shard @ {group.param_dtype}: " diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4e78b822..208e0279 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1503,6 +1503,17 @@ def set_default_diffusion_args(args) -> None: else: args.ref_mode = "none" + # --precision-default-dtype fills every dtype knob left unset; specific flags win. + if args.diffusion_forward_dtype is None: + args.diffusion_forward_dtype = args.precision_default_dtype or "bf16" + if args.precision_default_dtype is not None: + from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig + + # Mirrors the engine's forwarding rule: a value equal to the class default counts as unset. + sglang_dit = getattr(args, "sglang_dit_precision", None) + if sglang_dit is None or sglang_dit == getattr(PipelineConfig, "dit_precision", None): + args.sglang_dit_precision = args.precision_default_dtype + def miles_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) @@ -1530,20 +1541,11 @@ def miles_validate_args(args): if args.eval_reward_key is None: args.eval_reward_key = args.reward_key - # Resolve the precision default-dtype cascade: specific flag > --precision-default-dtype > built-in. - if args.diffusion_forward_dtype is None: - args.diffusion_forward_dtype = args.precision_default_dtype or "bf16" - if args.precision_default_dtype is not None: - from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig - - sglang_dit = getattr(args, "sglang_dit_precision", None) - if sglang_dit is None or sglang_dit == getattr(PipelineConfig, "dit_precision", None): - args.sglang_dit_precision = args.precision_default_dtype - elif sglang_dit != args.precision_default_dtype: - raise ValueError( - f"--sglang-dit-precision {sglang_dit} conflicts with " - f"--precision-default-dtype {args.precision_default_dtype}" - ) + if args.precision_default_dtype is not None and args.sglang_dit_precision != args.precision_default_dtype: + raise ValueError( + f"--sglang-dit-precision {args.sglang_dit_precision} conflicts with " + f"--precision-default-dtype {args.precision_default_dtype}" + ) args.update_weight_target_modules = [ name.strip() for name in args.update_weight_target_module.split(",") if name.strip() diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 555eed56..c315cff0 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -48,6 +48,7 @@ def test_module_cls_rule_lowers_norms_to_one_subshard_group(): assert group.param_dtype is torch.float32 assert group.modules == [model.blocks[0].norm, model.blocks[1].norm] assert set(group.param_fqns) == {f"blocks.{i}.norm.{n}" for i in range(2) for n in ("weight", "bias")} + assert all(cast.dtype is torch.float32 for cast in compiled.master_casts) compiled.apply_master_casts() assert model.blocks[0].norm.weight.dtype is torch.float32 assert model.blocks[0].linear.weight.dtype is torch.bfloat16 @@ -65,7 +66,7 @@ def test_param_fqn_rule_casts_buffer_master(): model = _model() spec = PrecisionSpec(rules=(Rule(ParamSel("*.freqs"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert {p.fqn for p in compiled.master_casts} == {"blocks.0.freqs", "blocks.1.freqs"} + assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.freqs", "blocks.1.freqs"} assert compiled.subshard_groups == [] compiled.apply_master_casts() assert model.blocks[0].freqs.dtype is torch.float32 From 926b9a310ade802d04a592692c194e0d3e562918 Mon Sep 17 00:00:00 2001 From: rockdu Date: Sun, 2 Aug 2026 21:30:10 -0700 Subject: [PATCH 05/50] docs(fsdp): comment each stage of compile_precision_plan --- miles/backends/fsdp_utils/precision.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index a8b5c37e..b344604e 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -137,6 +137,8 @@ def compile_precision_plan( for rule in spec.rules: _validate_rule(rule) + # Enumerate every param/buffer once: (fqn, tensor, is_buffer, owning module) plus + # each module's float-param set (needed later for whole-module coverage checks). named: list[tuple[str, torch.Tensor, bool, str, torch.nn.Module]] = [] float_params_of_module: dict[str, set[str]] = {} for mod_fqn, module in model.named_modules(): @@ -148,6 +150,8 @@ def compile_precision_plan( for name, buf in module.named_buffers(recurse=False): named.append((f"{prefix}{name}", buf, True, mod_fqn, module)) + # Precompile each rule into a tensor-fqn -> bool matcher; a ModuleSel matches + # every tensor under its matched modules' subtrees. matchers = [] for rule in spec.rules: if isinstance(rule.select, ParamSel): @@ -162,6 +166,9 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: return None return default_dtype if axis == "default" else _DTYPES[axis] + # Per tensor: fold matching rules in declaration order (last match wins per axis), + # then route each axis to its lowering: master -> load-time cast, gather -> sub-wrap + # request on the owning module (gather == default needs nothing, the block policy is it). match_counts = [0] * len(spec.rules) master_casts: list[MasterCast] = [] gather_by_module: dict[str, tuple[torch.nn.Module, dict[str, torch.dtype]]] = {} @@ -186,12 +193,14 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: if master_dtype is not None and master_dtype != tensor.dtype: master_casts.append(MasterCast(fqn, tensor, master_dtype)) + # A rule that matched nothing is almost certainly a typo'd pattern or class name. for rule, count in zip(spec.rules, match_counts, strict=True): if count == 0: raise ValueError(f"precision rule matched no tensor: {rule}") - # Gather overrides lower to whole-module nested fully_shard units, so every - # float param of an affected module must agree on one gather dtype. + # Lower gather requests to nested fully_shard units: fully_shard wraps whole modules, + # so an affected module's float params must all agree on one gather dtype, and modules + # sharing a dtype merge into one group (one extra all-gather per group per step). groups: dict[torch.dtype, SubShardGroup] = {} for mod_fqn, (module, requests) in gather_by_module.items(): dtypes = set(requests.values()) From ba941fd53c30fa79a66d58513687424b39e3e6c1 Mon Sep 17 00:00:00 2001 From: rockdu Date: Sun, 2 Aug 2026 23:25:15 -0700 Subject: [PATCH 06/50] docs(fsdp): replace stage comments with a flow diagram in the precision module docstring --- miles/backends/fsdp_utils/precision.py | 57 +++++++++++++++----------- 1 file changed, 32 insertions(+), 25 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index b344604e..e3c22083 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -1,23 +1,36 @@ """Fine-grained weight-precision control for FSDP2. A family declares per-tensor dtype intent as PrecisionSpec rules on its -TrainPipelineConfig (last matching rule wins per axis): +TrainPipelineConfig; each rule pins one or both axes (last match wins per axis): - master: resident dtype of the param/buffer (optimizer state precision) - gather: dtype the param is cast to for FSDP all-gather + forward -Compute dtype is deliberately not managed here: the trainer runs the DiT -forward under torch.autocast(default dtype); op-level exceptions belong to -the monkey-patch registry. - -``compile_precision_plan`` resolves rules to a per-tensor plan and lowers it -onto what FSDP2 can express: - - inline (default): tensor follows its wrap unit's MixedPrecisionPolicy - - sub-shard: gather pinned away from the default -> the owning modules are - grouped into one nested fully_shard unit per gather dtype with its own - MixedPrecisionPolicy, staying fully inside FSDP (DTensor params, FSDP - grad reduction, DCP/offload as usual) at the cost of one extra - all-gather per group per forward +Compute dtype is not managed here: the trainer autocasts the DiT forward and +op-level exceptions belong to the monkey-patch registry. + +``compile_precision_plan`` lowers the rules onto what FSDP2 can express: + + PrecisionSpec rules + | + (1) inventory every param/buffer: FQN, owning module + | + (2) per tensor, fold matching rules -> (master, gather) intent + | | + master axis gather axis + | | + (3) != loaded dtype? (4) != default dtype? --no--> inline: the block + -> MasterCast, | policy already is + cast at load time v the default + sub-wrap request on the owning module + | + (5) per module: one gather dtype + all float params covered, then + merge same-dtype modules -> one SubShardGroup per dtype + | + apply_fsdp2: fully_shard(group.modules, param_dtype=group dtype), + nested before the block/root wrap, one extra all-gather per group + Gather intent that does not align to a module boundary (e.g. a bare -nn.Parameter on a block) cannot be sub-wrapped and is rejected. +nn.Parameter on a block) cannot be sub-wrapped and is rejected; FSDP2 itself +additionally requires uniform master dtype among trainable params per unit. """ from __future__ import annotations @@ -137,8 +150,7 @@ def compile_precision_plan( for rule in spec.rules: _validate_rule(rule) - # Enumerate every param/buffer once: (fqn, tensor, is_buffer, owning module) plus - # each module's float-param set (needed later for whole-module coverage checks). + # (1) Tensor inventory, plus per-module float-param sets for the stage (5) coverage check. named: list[tuple[str, torch.Tensor, bool, str, torch.nn.Module]] = [] float_params_of_module: dict[str, set[str]] = {} for mod_fqn, module in model.named_modules(): @@ -150,8 +162,7 @@ def compile_precision_plan( for name, buf in module.named_buffers(recurse=False): named.append((f"{prefix}{name}", buf, True, mod_fqn, module)) - # Precompile each rule into a tensor-fqn -> bool matcher; a ModuleSel matches - # every tensor under its matched modules' subtrees. + # (2a) One fqn -> bool matcher per rule; ModuleSel covers its matched modules' subtrees. matchers = [] for rule in spec.rules: if isinstance(rule.select, ParamSel): @@ -166,9 +177,7 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: return None return default_dtype if axis == "default" else _DTYPES[axis] - # Per tensor: fold matching rules in declaration order (last match wins per axis), - # then route each axis to its lowering: master -> load-time cast, gather -> sub-wrap - # request on the owning module (gather == default needs nothing, the block policy is it). + # (2b)-(4) Fold rules per tensor, route master to casts and gather to sub-wrap requests. match_counts = [0] * len(spec.rules) master_casts: list[MasterCast] = [] gather_by_module: dict[str, tuple[torch.nn.Module, dict[str, torch.dtype]]] = {} @@ -193,14 +202,12 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: if master_dtype is not None and master_dtype != tensor.dtype: master_casts.append(MasterCast(fqn, tensor, master_dtype)) - # A rule that matched nothing is almost certainly a typo'd pattern or class name. + # A rule matching nothing is almost certainly a typo'd pattern or class name. for rule, count in zip(spec.rules, match_counts, strict=True): if count == 0: raise ValueError(f"precision rule matched no tensor: {rule}") - # Lower gather requests to nested fully_shard units: fully_shard wraps whole modules, - # so an affected module's float params must all agree on one gather dtype, and modules - # sharing a dtype merge into one group (one extra all-gather per group per step). + # (5) Validate module-boundary alignment, then merge same-dtype modules into groups. groups: dict[torch.dtype, SubShardGroup] = {} for mod_fqn, (module, requests) in gather_by_module.items(): dtypes = set(requests.values()) From bb6220b8b552ceabe98e9aeb293c0e5d9258a062 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 01:48:07 -0700 Subject: [PATCH 07/50] style: apply pre-commit isort to precision plan test imports --- tests/fast/backends/fsdp_utils/test_precision_plan.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index c315cff0..56af619b 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -6,13 +6,7 @@ import torch import torch.nn as nn -from miles.backends.fsdp_utils.precision import ( - ModuleSel, - ParamSel, - PrecisionSpec, - Rule, - compile_precision_plan, -) +from miles.backends.fsdp_utils.precision import ModuleSel, ParamSel, PrecisionSpec, Rule, compile_precision_plan class Block(nn.Module): From 0ca24cf9cd5dffa0c52abfcd94cd19a4bab8c37b Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 12:11:45 -0700 Subject: [PATCH 08/50] feat(fsdp): declare model-boundary input dtypes as family policy --- miles/backends/fsdp_utils/actor.py | 22 ++++-- .../configs/train_pipeline_config.py | 3 + miles/backends/fsdp_utils/precision.py | 43 +++++++++++ .../fsdp_utils/test_input_dtype_policy.py | 75 +++++++++++++++++++ 4 files changed, 136 insertions(+), 7 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_input_dtype_policy.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index e64408fd..9bbcc27e 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -39,7 +39,7 @@ from .lr_scheduler import get_lr_scheduler from .metrics import new_metric_buffer from .parallel import create_fsdp_parallel_state -from .precision import compile_precision_plan, log_precision_summary, resolve_dtype +from .precision import apply_input_dtype_policy, compile_precision_plan, log_precision_summary, resolve_dtype from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -524,16 +524,24 @@ def _forward_train_pair_batch( train_pipeline_config = self.train_pipeline_config forward_dtype = self._forward_dtype + # Boundary dtypes are family policy; op interiors stay autocast-managed. + latents_in, timesteps_in, (pos_cond_in, neg_cond_in, joint_cond_in) = apply_input_dtype_policy( + train_pipeline_config.input_dtype_policy, + latents=prepared.latents, + timesteps=prepared.timesteps_for_model, + conds=(prepared.pos_cond, prepared.neg_cond, prepared.joint_cond), + default_dtype=forward_dtype, + ) + def _compute_noise_pred() -> torch.Tensor: - # Inputs/params keep their resident dtypes; compute dtype is autocast-managed. with torch.autocast("cuda", dtype=forward_dtype, enabled=forward_dtype != torch.float32): return train_pipeline_config.compute_noise_pred( model=prepared.model, - latents_input=prepared.latents, - timesteps_input=prepared.timesteps_for_model, - pos_cond=prepared.pos_cond, - neg_cond=prepared.neg_cond, - joint_cond=prepared.joint_cond, + latents_input=latents_in, + timesteps_input=timesteps_in, + pos_cond=pos_cond_in, + neg_cond=neg_cond_in, + joint_cond=joint_cond_in, use_cfg=prepared.use_cfg, cfg_batching=prepared.cfg_batching, guidance_scale=prepared.guidance_scale, diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 049474e2..c63efffb 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -85,6 +85,9 @@ class TrainPipelineConfig(abc.ABC): rollout_patch_group: str | None = None # Weight-precision rules (master/gather dtypes) compiled onto FSDP2; see precision.py. precision_spec: PrecisionSpec = PrecisionSpec() + # Model-boundary input dtypes (see precision.apply_input_dtype_policy); the default + # mirrors the wan/ltx rollout boundary: latents/cond cast once, timestep kept fp32. + input_dtype_policy: dict = {"latents": "default", "cond": "default", "timestep": "fp32"} # Default component paths (miles custom-function style); CLI args override. model_backend_path: str = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" # Native model package import path; required when model_backend_path is MilesModelBackend. diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index e3c22083..63afa002 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -229,6 +229,49 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: return CompiledPrecision(master_casts=master_casts, subshard_groups=list(groups.values())) +# --------------------------------------------------------------------------- +# Boundary-input dtype policy (TrainPipelineConfig.input_dtype_policy) +# --------------------------------------------------------------------------- + +INPUT_DTYPE_POLICY_KEYS = ("latents", "cond", "timestep") + + +def apply_input_dtype_policy( + policy: dict, + *, + latents: torch.Tensor, + timesteps: torch.Tensor, + conds: tuple, + default_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, tuple]: + """Cast model-boundary inputs once, mirroring sglang-d's denoising stage + (autocast covers only matmul/conv ops, so e.g. a raw fp32 input reaching RoPE + would diverge from rollout). Axis values: "default" (the run's forward dtype), + a dtype name, or None (pass through); only floating tensors are cast.""" + unknown = set(policy) - set(INPUT_DTYPE_POLICY_KEYS) + if unknown: + raise ValueError(f"input_dtype_policy has unknown keys {sorted(unknown)}; known: {INPUT_DTYPE_POLICY_KEYS}") + + def _axis(key: str) -> torch.dtype | None: + axis = policy.get(key) + if axis is None: + return None + if axis != "default" and axis not in _DTYPES: + raise ValueError(f"input_dtype_policy[{key!r}] has unknown dtype {axis!r}") + return default_dtype if axis == "default" else _DTYPES[axis] + + def _cast(value, dtype: torch.dtype | None): + if dtype is None or not torch.is_tensor(value) or not value.is_floating_point(): + return value + return value.to(dtype) + + latents_dtype, timestep_dtype, cond_dtype = _axis("latents"), _axis("timestep"), _axis("cond") + out_conds = tuple( + None if cond is None else {key: _cast(value, cond_dtype) for key, value in cond.items()} for cond in conds + ) + return _cast(latents, latents_dtype), _cast(timesteps, timestep_dtype), out_conds + + def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: logger.info( f"precision[{component}]: default gather dtype {default_dtype}, " diff --git a/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py new file mode 100644 index 00000000..77f39f73 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py @@ -0,0 +1,75 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[]) + +import pytest +import torch + +from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig +from miles.backends.fsdp_utils.precision import apply_input_dtype_policy + +DEFAULT_POLICY = TrainPipelineConfig.input_dtype_policy + + +def _inputs(): + latents = torch.zeros(1, 4, 8, dtype=torch.float32) + timesteps = torch.tensor([857.69], dtype=torch.float32) + pos_cond = { + "context": torch.zeros(1, 3, 8, dtype=torch.float32), + "context_mask": torch.ones(1, 3, dtype=torch.int64), + } + return latents, timesteps, pos_cond + + +def test_default_policy_matches_rollout_boundary(): + latents, timesteps, pos_cond = _inputs() + out_latents, out_timesteps, (out_pos, out_neg, out_joint) = apply_input_dtype_policy( + DEFAULT_POLICY, + latents=latents, + timesteps=timesteps, + conds=(pos_cond, None, None), + default_dtype=torch.bfloat16, + ) + assert out_latents.dtype == torch.bfloat16 + assert out_timesteps.dtype == torch.float32 + assert out_pos["context"].dtype == torch.bfloat16 + assert out_pos["context_mask"].dtype == torch.int64 + assert out_neg is None and out_joint is None + + +def test_family_override_timestep_default(): + latents, timesteps, pos_cond = _inputs() + policy = {**DEFAULT_POLICY, "timestep": "default"} + _, out_timesteps, _ = apply_input_dtype_policy( + policy, + latents=latents, + timesteps=timesteps, + conds=(pos_cond, None, None), + default_dtype=torch.bfloat16, + ) + assert out_timesteps.dtype == torch.bfloat16 + + +def test_none_axis_passes_through(): + latents, timesteps, pos_cond = _inputs() + out_latents, _, (out_pos, _, _) = apply_input_dtype_policy( + {"latents": None, "cond": None, "timestep": None}, + latents=latents, + timesteps=timesteps, + conds=(pos_cond, None, None), + default_dtype=torch.bfloat16, + ) + assert out_latents.dtype == torch.float32 + assert out_pos["context"].dtype == torch.float32 + + +def test_unknown_key_rejected(): + latents, timesteps, pos_cond = _inputs() + with pytest.raises(ValueError, match="unknown keys"): + apply_input_dtype_policy( + {"latnets": "default"}, + latents=latents, + timesteps=timesteps, + conds=(pos_cond, None, None), + default_dtype=torch.bfloat16, + ) From 1214240fde0300e37f305efe3db612b93c0b4e35 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 15:58:30 -0700 Subject: [PATCH 09/50] test(e2e): update SD3.5 OCR standard for the autocast forward numerics --- .../test_sd3_ocr_grpo_2xGPU.json | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/ci/fixtures/e2e_standards/test_sd3_ocr_grpo_2xGPU.json b/tests/ci/fixtures/e2e_standards/test_sd3_ocr_grpo_2xGPU.json index 8c5771d4..44adde37 100644 --- a/tests/ci/fixtures/e2e_standards/test_sd3_ocr_grpo_2xGPU.json +++ b/tests/ci/fixtures/e2e_standards/test_sd3_ocr_grpo_2xGPU.json @@ -1,6 +1,6 @@ { "meta": { - "commit": "d4a7b1df1bad4a293dc70cc3b47f259767bf25ec", + "commit": "0ca24cf9cd5dffa0c52abfcd94cd19a4bab8c37b", "source": "test_sd3_ocr_grpo_2xGPU.py" }, "metrics": { @@ -11,7 +11,7 @@ ], [ 1, - 0.46756434440612793 + 0.49097996950149536 ] ], "rollout/reward/raw_median": [ @@ -21,7 +21,7 @@ ], [ 1, - 0.5 + 0.5333333611488342 ] ], "rollout/reward/raw_num_samples": [ @@ -41,61 +41,61 @@ ], [ 1, - 0.3667338788509369 + 0.3622446358203888 ] ], "train/grad_norm": [ [ 1.0, - 0.0018781634280458093 + 0.001604570308700204 ], [ 2.0, - 0.006335652898997068 + 0.003820131067186594 ], [ 3.0, - 0.002990065375342965 + 0.0013580435188487172 ], [ 4.0, - 0.004158319905400276 + 0.004337321501225233 ] ], "train/log_prob_mean_abs_diff": [ [ 1.0, - 0.0001490480382926762 + 9.299978846684098e-05 ], [ 2.0, - 0.0012502310797572135 + 0.0011721035931259393 ], [ 3.0, - 0.00013852929696440698 + 0.00011219287989661098 ], [ 4.0, - 0.0008732350077480077 + 0.001000555045902729 ] ], "train/log_prob_new_idx_0": [ [ 1.0, - -0.720925610512495 + -0.7209111824631691 ], [ 2.0, - -0.7189844809472561 + -0.7189970053732395 ], [ 3.0, - -0.720899011939764 + -0.720885843038559 ], [ 4.0, - -0.7225109003484249 + -0.7225540168583393 ] ], "train/log_prob_old_idx_0": [ From 2940cb71b51e47e97d65e2a0a209c5591bbd067b Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 17:03:41 -0700 Subject: [PATCH 10/50] refactor(fsdp): make precision rules module-granular, dropping ParamSel and boundary validation --- miles/backends/fsdp_utils/precision.py | 146 +++++++----------- .../fsdp_utils/test_precision_plan.py | 68 ++++---- 2 files changed, 86 insertions(+), 128 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 63afa002..7278a374 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -1,36 +1,36 @@ -"""Fine-grained weight-precision control for FSDP2. +"""Fine-grained weight-precision control for FSDP2, at module granularity. -A family declares per-tensor dtype intent as PrecisionSpec rules on its +A family declares per-module dtype intent as PrecisionSpec rules on its TrainPipelineConfig; each rule pins one or both axes (last match wins per axis): - - master: resident dtype of the param/buffer (optimizer state precision) - - gather: dtype the param is cast to for FSDP all-gather + forward -Compute dtype is not managed here: the trainer autocasts the DiT forward and -op-level exceptions belong to the monkey-patch registry. + - master: resident dtype of the module's params/buffers (optimizer precision) + - gather: dtype the params are cast to for FSDP all-gather + forward +Compute dtype is not managed here: the trainer autocasts the DiT forward +(boundary inputs via apply_input_dtype_policy below) and op-level exceptions +belong to the monkey-patch registry. ``compile_precision_plan`` lowers the rules onto what FSDP2 can express: PrecisionSpec rules | - (1) inventory every param/buffer: FQN, owning module + (1) match each rule to modules; a rule covers its modules' whole subtrees | - (2) per tensor, fold matching rules -> (master, gather) intent + (2) per module, fold covering rules -> (master, gather) intent | | master axis gather axis | | (3) != loaded dtype? (4) != default dtype? --no--> inline: the block - -> MasterCast, | policy already is - cast at load time v the default - sub-wrap request on the owning module - | - (5) per module: one gather dtype + all float params covered, then - merge same-dtype modules -> one SubShardGroup per dtype + -> MasterCast of the | policy already is + module's own float v the default + tensors, at load time sub-wrap: merge same-dtype modules into one + SubShardGroup per dtype (buffer-only and + paramless modules have nothing to gather: skipped) | apply_fsdp2: fully_shard(group.modules, param_dtype=group dtype), nested before the block/root wrap, one extra all-gather per group -Gather intent that does not align to a module boundary (e.g. a bare -nn.Parameter on a block) cannot be sub-wrapped and is rejected; FSDP2 itself -additionally requires uniform master dtype among trainable params per unit. +Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, +and FSDP2 requires uniform master dtype among trainable params per unit), so +finer-grained selectors are deliberately not offered. """ from __future__ import annotations @@ -57,24 +57,17 @@ def resolve_dtype(name: str) -> torch.dtype: @dataclass(frozen=True) class ModuleSel: - """Matches every param/buffer under modules with fnmatch-ing FQN or exact class name.""" + """Matches modules (and their subtrees) by fnmatch-ing FQN and/or exact class name.""" fqn: str | None = None cls: str | None = None -@dataclass(frozen=True) -class ParamSel: - """Matches individual params/buffers by fnmatch on their FQN.""" - - fqn: str - - @dataclass(frozen=True) class Rule: """Axes take a dtype name ("fp32"/"bf16"/"fp16"), "default" (the run's default dtype), or None (untouched).""" - select: ModuleSel | ParamSel + select: ModuleSel master: str | None = None gather: str | None = None @@ -118,7 +111,7 @@ def apply_master_casts(self) -> None: def _validate_rule(rule: Rule) -> None: if rule.master is None and rule.gather is None: raise ValueError(f"precision rule sets no dtype axis: {rule}") - if isinstance(rule.select, ModuleSel) and rule.select.fqn is None and rule.select.cls is None: + if rule.select.fqn is None and rule.select.cls is None: raise ValueError(f"precision rule has an empty ModuleSel: {rule}") for axis in (rule.master, rule.gather): if axis is not None and axis != "default" and axis not in _DTYPES: @@ -142,89 +135,56 @@ def compile_precision_plan( *, default_dtype: torch.dtype, ) -> CompiledPrecision: - """Resolve spec rules against the (pre-LoRA, pre-FSDP) model and lower them. - - Raises on rules that match nothing (likely a typo'd pattern) and on plans - FSDP2 cannot express (see module docstring). - """ + """Resolve spec rules against the (pre-LoRA, pre-FSDP) model and lower them per module.""" for rule in spec.rules: _validate_rule(rule) - # (1) Tensor inventory, plus per-module float-param sets for the stage (5) coverage check. - named: list[tuple[str, torch.Tensor, bool, str, torch.nn.Module]] = [] - float_params_of_module: dict[str, set[str]] = {} - for mod_fqn, module in model.named_modules(): - prefix = f"{mod_fqn}." if mod_fqn else "" - for name, param in module.named_parameters(recurse=False): - named.append((f"{prefix}{name}", param, False, mod_fqn, module)) - if param.is_floating_point(): - float_params_of_module.setdefault(mod_fqn, set()).add(f"{prefix}{name}") - for name, buf in module.named_buffers(recurse=False): - named.append((f"{prefix}{name}", buf, True, mod_fqn, module)) - - # (2a) One fqn -> bool matcher per rule; ModuleSel covers its matched modules' subtrees. - matchers = [] - for rule in spec.rules: - if isinstance(rule.select, ParamSel): - pattern = rule.select.fqn - matchers.append(lambda fqn, pattern=pattern: fnmatch(fqn, pattern)) - else: - prefixes = _matched_module_prefixes(rule.select, model) - matchers.append(lambda fqn, prefixes=prefixes: any(p == "" or fqn.startswith(f"{p}.") for p in prefixes)) + # (1) Match each rule to modules once; a rule that matches nothing is almost certainly a typo. + rule_prefixes = [_matched_module_prefixes(rule.select, model) for rule in spec.rules] + for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): + if not prefixes: + raise ValueError(f"precision rule matched no module: {rule}") + + def _covers(prefixes: list[str], mod_fqn: str) -> bool: + return any(p == "" or mod_fqn == p or mod_fqn.startswith(f"{p}.") for p in prefixes) def _resolve_axis(axis: str | None) -> torch.dtype | None: if axis is None: return None return default_dtype if axis == "default" else _DTYPES[axis] - # (2b)-(4) Fold rules per tensor, route master to casts and gather to sub-wrap requests. - match_counts = [0] * len(spec.rules) + # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard groups. master_casts: list[MasterCast] = [] - gather_by_module: dict[str, tuple[torch.nn.Module, dict[str, torch.dtype]]] = {} - for fqn, tensor, is_buffer, mod_fqn, module in named: + groups: dict[torch.dtype, SubShardGroup] = {} + for mod_fqn, module in model.named_modules(): master = gather = None - for i, rule in enumerate(spec.rules): - if not matchers[i](fqn): + for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): + if not _covers(prefixes, mod_fqn): continue - match_counts[i] += 1 master = rule.master if rule.master is not None else master gather = rule.gather if rule.gather is not None else gather + if master is None and gather is None: + continue - gather_dtype = _resolve_axis(gather) - if gather_dtype is not None: - if is_buffer: - raise ValueError(f"precision rule pins gather dtype on buffer {fqn}; buffers are never gathered") - if gather_dtype != default_dtype: - if mod_fqn == "": - raise ValueError(f"{fqn}: cannot sub-wrap the root module for a gather override") - gather_by_module.setdefault(mod_fqn, (module, {}))[1][fqn] = gather_dtype + prefix = f"{mod_fqn}." if mod_fqn else "" + params = list(module.named_parameters(recurse=False)) + tensors = params + list(module.named_buffers(recurse=False)) master_dtype = _resolve_axis(master) - if master_dtype is not None and master_dtype != tensor.dtype: - master_casts.append(MasterCast(fqn, tensor, master_dtype)) + if master_dtype is not None: + for name, tensor in tensors: + if tensor.is_floating_point() and tensor.dtype != master_dtype: + master_casts.append(MasterCast(f"{prefix}{name}", tensor, master_dtype)) - # A rule matching nothing is almost certainly a typo'd pattern or class name. - for rule, count in zip(spec.rules, match_counts, strict=True): - if count == 0: - raise ValueError(f"precision rule matched no tensor: {rule}") - - # (5) Validate module-boundary alignment, then merge same-dtype modules into groups. - groups: dict[torch.dtype, SubShardGroup] = {} - for mod_fqn, (module, requests) in gather_by_module.items(): - dtypes = set(requests.values()) - if len(dtypes) > 1: - raise ValueError( - f"module {mod_fqn} mixes gather dtypes {sorted(map(str, dtypes))}; FSDP sub-wrap needs one" - ) - missing = float_params_of_module.get(mod_fqn, set()) - requests.keys() - if missing: - raise ValueError( - f"gather override must cover the whole module {mod_fqn} for FSDP sub-wrap; " - f"missing {sorted(missing)}" - ) - dtype = dtypes.pop() - group = groups.setdefault(dtype, SubShardGroup(param_dtype=dtype)) - group.modules.append(module) - group.param_fqns.extend(sorted(requests)) + gather_dtype = _resolve_axis(gather) + if gather_dtype is not None and gather_dtype != default_dtype: + float_params = [name for name, param in params if param.is_floating_point()] + if not float_params: + continue + if mod_fqn == "": + raise ValueError("cannot sub-wrap the root module for a gather override") + group = groups.setdefault(gather_dtype, SubShardGroup(param_dtype=gather_dtype)) + group.modules.append(module) + group.param_fqns.extend(f"{prefix}{name}" for name in float_params) return CompiledPrecision(master_casts=master_casts, subshard_groups=list(groups.values())) diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 56af619b..883d889c 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -6,7 +6,13 @@ import torch import torch.nn as nn -from miles.backends.fsdp_utils.precision import ModuleSel, ParamSel, PrecisionSpec, Rule, compile_precision_plan +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision_plan + + +class Rope(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("freqs", torch.zeros(4)) class Block(nn.Module): @@ -14,7 +20,7 @@ def __init__(self): super().__init__() self.linear = nn.Linear(8, 8) self.norm = nn.LayerNorm(8) - self.register_buffer("freqs", torch.zeros(4)) + self.rope = Rope() class Tiny(nn.Module): @@ -33,7 +39,7 @@ def test_empty_spec_compiles_to_nothing(): assert compiled.subshard_groups == [] -def test_module_cls_rule_lowers_norms_to_one_subshard_group(): +def test_cls_rule_lowers_norms_to_one_subshard_group(): model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -42,7 +48,6 @@ def test_module_cls_rule_lowers_norms_to_one_subshard_group(): assert group.param_dtype is torch.float32 assert group.modules == [model.blocks[0].norm, model.blocks[1].norm] assert set(group.param_fqns) == {f"blocks.{i}.norm.{n}" for i in range(2) for n in ("weight", "bias")} - assert all(cast.dtype is torch.float32 for cast in compiled.master_casts) compiled.apply_master_casts() assert model.blocks[0].norm.weight.dtype is torch.float32 assert model.blocks[0].linear.weight.dtype is torch.bfloat16 @@ -56,44 +61,37 @@ def test_gather_without_master_is_allowed(): assert compiled.subshard_groups[0].param_dtype is torch.float32 -def test_param_fqn_rule_casts_buffer_master(): +def test_master_rule_covers_matched_subtree(): model = _model() - spec = PrecisionSpec(rules=(Rule(ParamSel("*.freqs"), master="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.freqs", "blocks.1.freqs"} - assert compiled.subshard_groups == [] + assert all(cast.fqn.startswith("blocks.1.") for cast in compiled.master_casts) compiled.apply_master_casts() - assert model.blocks[0].freqs.dtype is torch.float32 + assert model.blocks[1].linear.weight.dtype is torch.float32 + assert model.blocks[1].rope.freqs.dtype is torch.float32 + assert model.blocks[0].linear.weight.dtype is torch.bfloat16 -def test_gather_matching_default_is_inline(): - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) +def test_buffer_only_module_master_cast(): + model = _model() + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32"),)) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} assert compiled.subshard_groups == [] - assert compiled.master_casts == [] - - -def test_partial_module_gather_rejected(): - spec = PrecisionSpec(rules=(Rule(ParamSel("*.norm.weight"), gather="fp32"),)) - with pytest.raises(ValueError, match="whole module"): - compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) -def test_mixed_gather_dtypes_in_module_rejected(): - spec = PrecisionSpec( - rules=( - Rule(ParamSel("*.norm.weight"), gather="fp32"), - Rule(ParamSel("*.norm.bias"), gather="fp16"), - ) - ) - with pytest.raises(ValueError, match="mixes gather dtypes"): - compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) +def test_buffer_only_module_gather_is_skipped(): + # Buffers are never gathered; a gather pin on a paramless module lowers to nothing. + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), gather="fp32"),)) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert compiled.subshard_groups == [] -def test_gather_on_buffer_rejected(): - spec = PrecisionSpec(rules=(Rule(ParamSel("*.freqs"), master="fp32", gather="fp32"),)) - with pytest.raises(ValueError, match="buffer"): - compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) +def test_gather_matching_default_is_inline(): + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert compiled.subshard_groups == [] + assert compiled.master_casts == [] def test_last_rule_wins_per_axis(): @@ -101,7 +99,7 @@ def test_last_rule_wins_per_axis(): spec = PrecisionSpec( rules=( Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"), - Rule(ParamSel("blocks.0.norm.*"), master="default", gather="default"), + Rule(ModuleSel(fqn="blocks.0.norm"), master="default", gather="default"), ) ) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -109,6 +107,6 @@ def test_last_rule_wins_per_axis(): def test_unmatched_rule_rejected(): - spec = PrecisionSpec(rules=(Rule(ParamSel("no.such.param"), master="fp32"),)) - with pytest.raises(ValueError, match="matched no tensor"): + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) + with pytest.raises(ValueError, match="matched no module"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) From f1d0101ffd363714af819113667e061f64eaa0bc Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 17:14:47 -0700 Subject: [PATCH 11/50] docs(fsdp): align stale dtype comments with the boundary input policy --- miles/backends/fsdp_utils/actor.py | 4 ++-- .../backends/fsdp_utils/configs/train_pipeline_config.py | 2 +- miles/backends/fsdp_utils/loss_hub/flow_grpo.py | 2 +- miles/backends/fsdp_utils/loss_hub/nft.py | 2 +- miles/backends/fsdp_utils/precision.py | 8 ++++---- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 9bbcc27e..7fb5a115 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -643,8 +643,8 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules def _fsdp_kwargs(policy_param_dtype): return { - # Inputs are not cast: compute dtype comes from the trainer's autocast, which - # also keeps grad-ckpt recompute dtypes consistent with the forward. + # FSDP casts no inputs (boundary casts belong to input_dtype_policy); compute dtype + # comes from the trainer's autocast, keeping grad-ckpt recompute dtype-consistent. "mp_policy": MixedPrecisionPolicy( param_dtype=policy_param_dtype, reduce_dtype=reduce_dtype, diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index c63efffb..152aef4b 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -86,7 +86,7 @@ class TrainPipelineConfig(abc.ABC): # Weight-precision rules (master/gather dtypes) compiled onto FSDP2; see precision.py. precision_spec: PrecisionSpec = PrecisionSpec() # Model-boundary input dtypes (see precision.apply_input_dtype_policy); the default - # mirrors the wan/ltx rollout boundary: latents/cond cast once, timestep kept fp32. + # casts latents/cond once to the forward dtype and keeps the timestep exact in fp32. input_dtype_policy: dict = {"latents": "default", "cond": "default", "timestep": "fp32"} # Default component paths (miles custom-function style); CLI args override. model_backend_path: str = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 162cb7d2..876d69dc 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -73,7 +73,7 @@ def prepare_flow_grpo_batch( if use_cfg else None ) - # Cond tensors keep their rollout dtypes; the DiT forward runs under autocast. + # Cond dtypes are set at the model boundary by the family input_dtype_policy (see actor). cfg_batching = use_cfg and bool(args.fsdp_cfg_batching) joint_cond = pos_cond = neg_cond = None if cfg_batching: diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 51e2c678..97932078 100644 --- a/miles/backends/fsdp_utils/loss_hub/nft.py +++ b/miles/backends/fsdp_utils/loss_hub/nft.py @@ -37,7 +37,7 @@ def prepare_nft_batch( component_name, model = next(iter(ctx.models.items())) pos_list = [config.prepare_cond_kwargs(batch[i]["denoising_env"].pos_cond_kwargs, device) for i in range(bsz)] - # Cond tensors keep their rollout dtypes; the DiT forward runs under autocast. + # Cond dtypes are set at the model boundary by the family input_dtype_policy (see actor). pos_cond = config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len) num_train_timesteps = ctx.scheduler.config.num_train_timesteps diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 7278a374..0f1465bc 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -204,10 +204,10 @@ def apply_input_dtype_policy( conds: tuple, default_dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor, tuple]: - """Cast model-boundary inputs once, mirroring sglang-d's denoising stage - (autocast covers only matmul/conv ops, so e.g. a raw fp32 input reaching RoPE - would diverge from rollout). Axis values: "default" (the run's forward dtype), - a dtype name, or None (pass through); only floating tensors are cast.""" + """Cast model-boundary inputs once: autocast covers only matmul/conv ops, so + without a boundary cast e.g. a raw fp32 latent keeps fp32 through element-wise + ops. Axis values: "default" (the run's forward dtype), a dtype name, or None + (pass through); only floating tensors are cast.""" unknown = set(policy) - set(INPUT_DTYPE_POLICY_KEYS) if unknown: raise ValueError(f"input_dtype_policy has unknown keys {sorted(unknown)}; known: {INPUT_DTYPE_POLICY_KEYS}") From 90c7f688e380516ab5ef1cba8235a6739e19829c Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 17:16:10 -0700 Subject: [PATCH 12/50] docs(fsdp): correct the precision module docstring after the boundary policy and ZeRO-2 changes --- miles/backends/fsdp_utils/precision.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 0f1465bc..c8d1fc50 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -4,9 +4,10 @@ TrainPipelineConfig; each rule pins one or both axes (last match wins per axis): - master: resident dtype of the module's params/buffers (optimizer precision) - gather: dtype the params are cast to for FSDP all-gather + forward -Compute dtype is not managed here: the trainer autocasts the DiT forward -(boundary inputs via apply_input_dtype_policy below) and op-level exceptions -belong to the monkey-patch registry. +The weight spec does not manage compute dtype: the trainer autocasts the DiT +forward, model-boundary input dtypes are family policy applied by +``apply_input_dtype_policy`` below, and op-level exceptions belong to the +monkey-patch registry. ``compile_precision_plan`` lowers the rules onto what FSDP2 can express: @@ -26,7 +27,8 @@ paramless modules have nothing to gather: skipped) | apply_fsdp2: fully_shard(group.modules, param_dtype=group dtype), - nested before the block/root wrap, one extra all-gather per group + nested before the block/root wrap; one extra all-gather per group per + step (reshard_after_forward=False, so backward re-uses the forward gather) Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and FSDP2 requires uniform master dtype among trainable params per unit), so From ed0068858e1c99c3c3358f7df4ad4d3648110ec0 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 17:18:28 -0700 Subject: [PATCH 13/50] refactor(fsdp): default input_dtype_policy to passthrough --- .../fsdp_utils/configs/train_pipeline_config.py | 4 ++-- .../fsdp_utils/test_input_dtype_policy.py | 15 ++++++++------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 152aef4b..1ac2588b 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -86,8 +86,8 @@ class TrainPipelineConfig(abc.ABC): # Weight-precision rules (master/gather dtypes) compiled onto FSDP2; see precision.py. precision_spec: PrecisionSpec = PrecisionSpec() # Model-boundary input dtypes (see precision.apply_input_dtype_policy); the default - # casts latents/cond once to the forward dtype and keeps the timestep exact in fp32. - input_dtype_policy: dict = {"latents": "default", "cond": "default", "timestep": "fp32"} + # passes every input through unchanged, families opt into boundary casts explicitly. + input_dtype_policy: dict = {"latents": None, "cond": None, "timestep": None} # Default component paths (miles custom-function style); CLI args override. model_backend_path: str = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" # Native model package import path; required when model_backend_path is MilesModelBackend. diff --git a/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py index 77f39f73..facd63de 100644 --- a/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py +++ b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py @@ -21,7 +21,7 @@ def _inputs(): return latents, timesteps, pos_cond -def test_default_policy_matches_rollout_boundary(): +def test_default_policy_is_passthrough(): latents, timesteps, pos_cond = _inputs() out_latents, out_timesteps, (out_pos, out_neg, out_joint) = apply_input_dtype_policy( DEFAULT_POLICY, @@ -30,9 +30,9 @@ def test_default_policy_matches_rollout_boundary(): conds=(pos_cond, None, None), default_dtype=torch.bfloat16, ) - assert out_latents.dtype == torch.bfloat16 + assert out_latents.dtype == torch.float32 assert out_timesteps.dtype == torch.float32 - assert out_pos["context"].dtype == torch.bfloat16 + assert out_pos["context"].dtype == torch.float32 assert out_pos["context_mask"].dtype == torch.int64 assert out_neg is None and out_joint is None @@ -50,17 +50,18 @@ def test_family_override_timestep_default(): assert out_timesteps.dtype == torch.bfloat16 -def test_none_axis_passes_through(): +def test_cast_policy_casts_floats_only(): latents, timesteps, pos_cond = _inputs() out_latents, _, (out_pos, _, _) = apply_input_dtype_policy( - {"latents": None, "cond": None, "timestep": None}, + {"latents": "default", "cond": "default", "timestep": "fp32"}, latents=latents, timesteps=timesteps, conds=(pos_cond, None, None), default_dtype=torch.bfloat16, ) - assert out_latents.dtype == torch.float32 - assert out_pos["context"].dtype == torch.float32 + assert out_latents.dtype == torch.bfloat16 + assert out_pos["context"].dtype == torch.bfloat16 + assert out_pos["context_mask"].dtype == torch.int64 def test_unknown_key_rejected(): From 9239fd922869c9d0842f617b35d3f6bbec94342e Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 20:33:10 -0700 Subject: [PATCH 14/50] fix(fsdp): LTX opts into boundary input casts to keep rollout-aligned bf16 math --- miles/backends/fsdp_utils/configs/ltx.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index 91271f88..cf3b9df8 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -25,6 +25,9 @@ class LTXTrainPipelineConfig(TrainPipelineConfig): model_package = "miles.backends.fsdp_utils.models.ltx" # Audio branch has no optimizer state: we only train the video stream. optimizer_state_allowed_missing = ["audio"] + # forward_velocity anchors its element-wise math on latents.dtype and rollout runs + # it in bf16, so cast latents/cond at the boundary (fp32 passthrough drifts ~5x). + input_dtype_policy = {"latents": "default", "cond": "default", "timestep": None} def configure(self, args: Namespace) -> None: self._height = args.diffusion_height From f5ed76ae8ed45a15d54fbaed52bd58f5dfec7400 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 20:37:50 -0700 Subject: [PATCH 15/50] docs(fsdp): compress comments to one line per point --- miles/backends/fsdp_utils/actor.py | 3 +-- miles/backends/fsdp_utils/configs/ltx.py | 3 +-- .../backends/fsdp_utils/configs/train_pipeline_config.py | 3 +-- miles/backends/fsdp_utils/precision.py | 8 +++----- 4 files changed, 6 insertions(+), 11 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 7fb5a115..13999208 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -643,8 +643,7 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules def _fsdp_kwargs(policy_param_dtype): return { - # FSDP casts no inputs (boundary casts belong to input_dtype_policy); compute dtype - # comes from the trainer's autocast, keeping grad-ckpt recompute dtype-consistent. + # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent. "mp_policy": MixedPrecisionPolicy( param_dtype=policy_param_dtype, reduce_dtype=reduce_dtype, diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index cf3b9df8..f4a2604c 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -25,8 +25,7 @@ class LTXTrainPipelineConfig(TrainPipelineConfig): model_package = "miles.backends.fsdp_utils.models.ltx" # Audio branch has no optimizer state: we only train the video stream. optimizer_state_allowed_missing = ["audio"] - # forward_velocity anchors its element-wise math on latents.dtype and rollout runs - # it in bf16, so cast latents/cond at the boundary (fp32 passthrough drifts ~5x). + # forward_velocity anchors element-wise math on latents.dtype; rollout runs it bf16, so cast at the boundary. input_dtype_policy = {"latents": "default", "cond": "default", "timestep": None} def configure(self, args: Namespace) -> None: diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 1ac2588b..edace089 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -85,8 +85,7 @@ class TrainPipelineConfig(abc.ABC): rollout_patch_group: str | None = None # Weight-precision rules (master/gather dtypes) compiled onto FSDP2; see precision.py. precision_spec: PrecisionSpec = PrecisionSpec() - # Model-boundary input dtypes (see precision.apply_input_dtype_policy); the default - # passes every input through unchanged, families opt into boundary casts explicitly. + # Model-boundary input dtypes (see precision.apply_input_dtype_policy); families opt into casts explicitly. input_dtype_policy: dict = {"latents": None, "cond": None, "timestep": None} # Default component paths (miles custom-function style); CLI args override. model_backend_path: str = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index c8d1fc50..18db2f4e 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -80,7 +80,7 @@ class PrecisionSpec: # --------------------------------------------------------------------------- -# Compiler: per-tensor plan -> FSDP2 lowering (master casts + sub-shard groups) +# Compiler: per-module plan -> FSDP2 lowering (master casts + sub-shard groups) # --------------------------------------------------------------------------- @@ -206,10 +206,8 @@ def apply_input_dtype_policy( conds: tuple, default_dtype: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor, tuple]: - """Cast model-boundary inputs once: autocast covers only matmul/conv ops, so - without a boundary cast e.g. a raw fp32 latent keeps fp32 through element-wise - ops. Axis values: "default" (the run's forward dtype), a dtype name, or None - (pass through); only floating tensors are cast.""" + """Cast float boundary inputs per family policy ("default"/dtype name/None=passthrough); + autocast alone would leave element-wise ops running at the raw input dtype.""" unknown = set(policy) - set(INPUT_DTYPE_POLICY_KEYS) if unknown: raise ValueError(f"input_dtype_policy has unknown keys {sorted(unknown)}; known: {INPUT_DTYPE_POLICY_KEYS}") From 834b469e2a111b8548afa18f76d9746bf11cb8e5 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 20:48:16 -0700 Subject: [PATCH 16/50] refactor(fsdp): share one axis resolver between the compiler and the input policy --- miles/backends/fsdp_utils/precision.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 18db2f4e..6e9aa042 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -52,6 +52,13 @@ def resolve_dtype(name: str) -> torch.dtype: return _DTYPES[name] +def _resolve_axis(axis: str | None, default_dtype: torch.dtype) -> torch.dtype | None: + """Shared axis semantics: None -> untouched, "default" -> the run's default dtype, else a dtype name.""" + if axis is None: + return None + return default_dtype if axis == "default" else _DTYPES[axis] + + # --------------------------------------------------------------------------- # Spec: per-family declaration (see TrainPipelineConfig.precision_spec) # --------------------------------------------------------------------------- @@ -150,11 +157,6 @@ def compile_precision_plan( def _covers(prefixes: list[str], mod_fqn: str) -> bool: return any(p == "" or mod_fqn == p or mod_fqn.startswith(f"{p}.") for p in prefixes) - def _resolve_axis(axis: str | None) -> torch.dtype | None: - if axis is None: - return None - return default_dtype if axis == "default" else _DTYPES[axis] - # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard groups. master_casts: list[MasterCast] = [] groups: dict[torch.dtype, SubShardGroup] = {} @@ -171,13 +173,13 @@ def _resolve_axis(axis: str | None) -> torch.dtype | None: prefix = f"{mod_fqn}." if mod_fqn else "" params = list(module.named_parameters(recurse=False)) tensors = params + list(module.named_buffers(recurse=False)) - master_dtype = _resolve_axis(master) + master_dtype = _resolve_axis(master, default_dtype) if master_dtype is not None: for name, tensor in tensors: if tensor.is_floating_point() and tensor.dtype != master_dtype: master_casts.append(MasterCast(f"{prefix}{name}", tensor, master_dtype)) - gather_dtype = _resolve_axis(gather) + gather_dtype = _resolve_axis(gather, default_dtype) if gather_dtype is not None and gather_dtype != default_dtype: float_params = [name for name, param in params if param.is_floating_point()] if not float_params: @@ -214,11 +216,9 @@ def apply_input_dtype_policy( def _axis(key: str) -> torch.dtype | None: axis = policy.get(key) - if axis is None: - return None - if axis != "default" and axis not in _DTYPES: + if axis is not None and axis != "default" and axis not in _DTYPES: raise ValueError(f"input_dtype_policy[{key!r}] has unknown dtype {axis!r}") - return default_dtype if axis == "default" else _DTYPES[axis] + return _resolve_axis(axis, default_dtype) def _cast(value, dtype: torch.dtype | None): if dtype is None or not torch.is_tensor(value) or not value.is_floating_point(): From e987cdbe27cfab6446ff17ede81b72a895dbd6dd Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 21:02:13 -0700 Subject: [PATCH 17/50] refactor(fsdp): resolve precision rules on the module tree, nearest and most-specific annotation wins --- miles/backends/fsdp_utils/precision.py | 101 +++++++++++------- .../fsdp_utils/test_precision_plan.py | 67 ++++++++---- 2 files changed, 105 insertions(+), 63 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 6e9aa042..394c9143 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -1,11 +1,16 @@ """Fine-grained weight-precision control for FSDP2, at module granularity. A family declares per-module dtype intent as PrecisionSpec rules on its -TrainPipelineConfig; each rule pins one or both axes (last match wins per axis): +TrainPipelineConfig. A rule annotates the module-FQN tree nodes its ``path`` +matches (segment-wise glob, ``*`` never crosses dots) with one or both axes: - master: resident dtype of the module's params/buffers (optimizer precision) - gather: dtype the params are cast to for FSDP all-gather + forward -The weight spec does not manage compute dtype: the trainer autocasts the DiT -forward, model-boundary input dtypes are family policy applied by +Rule order carries no meaning; per axis each module resolves upward from itself: + - nearest wins: the closest annotated node on the self -> root chain applies + - most specific wins: on one node, the pattern with more literal segments + applies; a full tie is a compile error +Compute dtype is not managed here: the trainer autocasts the DiT forward, +model-boundary input dtypes are family policy applied by ``apply_input_dtype_policy`` below, and op-level exceptions belong to the monkey-patch registry. @@ -13,9 +18,9 @@ PrecisionSpec rules | - (1) match each rule to modules; a rule covers its modules' whole subtrees + (1) annotate matched tree nodes; a rule matching nothing is an error | - (2) per module, fold covering rules -> (master, gather) intent + (2) per module, resolve (master, gather) leaf-upward as above | | master axis gather axis | | @@ -64,19 +69,11 @@ def _resolve_axis(axis: str | None, default_dtype: torch.dtype) -> torch.dtype | # --------------------------------------------------------------------------- -@dataclass(frozen=True) -class ModuleSel: - """Matches modules (and their subtrees) by fnmatch-ing FQN and/or exact class name.""" - - fqn: str | None = None - cls: str | None = None - - @dataclass(frozen=True) class Rule: - """Axes take a dtype name ("fp32"/"bf16"/"fp16"), "default" (the run's default dtype), or None (untouched).""" + """path: segment-wise glob over module FQNs; axes take a dtype name, "default", or None (untouched).""" - select: ModuleSel + path: str master: str | None = None gather: str | None = None @@ -120,22 +117,36 @@ def apply_master_casts(self) -> None: def _validate_rule(rule: Rule) -> None: if rule.master is None and rule.gather is None: raise ValueError(f"precision rule sets no dtype axis: {rule}") - if rule.select.fqn is None and rule.select.cls is None: - raise ValueError(f"precision rule has an empty ModuleSel: {rule}") for axis in (rule.master, rule.gather): if axis is not None and axis != "default" and axis not in _DTYPES: raise ValueError(f"precision rule has unknown dtype {axis!r}: {rule}") -def _matched_module_prefixes(sel: ModuleSel, model: torch.nn.Module) -> list[str]: - prefixes = [] - for mod_fqn, module in model.named_modules(): - if sel.cls is not None and type(module).__name__ != sel.cls: - continue - if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): - continue - prefixes.append(mod_fqn) - return prefixes +def _path_matches(path: str, fqn: str) -> bool: + if path == "" or fqn == "": + return path == fqn + pattern_segments = path.split(".") + fqn_segments = fqn.split(".") + return len(pattern_segments) == len(fqn_segments) and all( + fnmatch(seg, pat) for seg, pat in zip(fqn_segments, pattern_segments, strict=True) + ) + + +def _specificity(path: str) -> int: + return sum(1 for seg in path.split(".") if not any(c in seg for c in "*?[")) + + +def _resolve_node_axis(rules_at_node: list[Rule], axis_name: str, mod_fqn: str) -> str | None: + """Most-specific-wins on one node; a full tie between distinct values is a compile error.""" + setters = [rule for rule in rules_at_node if getattr(rule, axis_name) is not None] + if not setters: + return None + best = max(_specificity(rule.path) for rule in setters) + winners = [rule for rule in setters if _specificity(rule.path) == best] + values = {getattr(rule, axis_name) for rule in winners} + if len(values) > 1: + raise ValueError(f"precision rules tie on {mod_fqn}.{axis_name}: {winners}") + return values.pop() def compile_precision_plan( @@ -148,25 +159,33 @@ def compile_precision_plan( for rule in spec.rules: _validate_rule(rule) - # (1) Match each rule to modules once; a rule that matches nothing is almost certainly a typo. - rule_prefixes = [_matched_module_prefixes(rule.select, model) for rule in spec.rules] - for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): - if not prefixes: + # (1) Annotate matched tree nodes; a rule that matches nothing is almost certainly a typo. + all_fqns = [fqn for fqn, _ in model.named_modules()] + annotations: dict[str, list[Rule]] = {} + for rule in spec.rules: + matched = [fqn for fqn in all_fqns if _path_matches(rule.path, fqn)] + if not matched: raise ValueError(f"precision rule matched no module: {rule}") - - def _covers(prefixes: list[str], mod_fqn: str) -> bool: - return any(p == "" or mod_fqn == p or mod_fqn.startswith(f"{p}.") for p in prefixes) - - # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard groups. + for fqn in matched: + annotations.setdefault(fqn, []).append(rule) + + def _resolve(mod_fqn: str, axis_name: str) -> str | None: + node = mod_fqn + while True: + if node in annotations: + value = _resolve_node_axis(annotations[node], axis_name, node) + if value is not None: + return value + if node == "": + return None + node = node.rsplit(".", 1)[0] if "." in node else "" + + # (2)-(4) Resolve each module leaf-upward, lower master to casts and gather to sub-shard groups. master_casts: list[MasterCast] = [] groups: dict[torch.dtype, SubShardGroup] = {} for mod_fqn, module in model.named_modules(): - master = gather = None - for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): - if not _covers(prefixes, mod_fqn): - continue - master = rule.master if rule.master is not None else master - gather = rule.gather if rule.gather is not None else gather + master = _resolve(mod_fqn, "master") + gather = _resolve(mod_fqn, "gather") if master is None and gather is None: continue diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 883d889c..8a48e0f6 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn -from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision_plan +from miles.backends.fsdp_utils.precision import PrecisionSpec, Rule, compile_precision_plan class Rope(nn.Module): @@ -39,9 +39,9 @@ def test_empty_spec_compiles_to_nothing(): assert compiled.subshard_groups == [] -def test_cls_rule_lowers_norms_to_one_subshard_group(): +def test_wildcard_rule_lowers_norms_to_one_subshard_group(): model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule("blocks.*.norm", master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert len(compiled.subshard_groups) == 1 group = compiled.subshard_groups[0] @@ -55,15 +55,15 @@ def test_cls_rule_lowers_norms_to_one_subshard_group(): def test_gather_without_master_is_allowed(): # Nested policy casts at all-gather, so gather may diverge from master. - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule("blocks.*.norm", gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert compiled.subshard_groups[0].param_dtype is torch.float32 -def test_master_rule_covers_matched_subtree(): +def test_ancestor_rule_covers_subtree(): model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) + spec = PrecisionSpec(rules=(Rule("blocks.1", master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert all(cast.fqn.startswith("blocks.1.") for cast in compiled.master_casts) compiled.apply_master_casts() @@ -72,9 +72,44 @@ def test_master_rule_covers_matched_subtree(): assert model.blocks[0].linear.weight.dtype is torch.bfloat16 +def test_nearest_annotation_wins(): + model = _model() + spec = PrecisionSpec( + rules=( + Rule("blocks.0.norm", master="default"), + Rule("blocks.0", master="fp32"), + ) + ) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + compiled.apply_master_casts() + assert model.blocks[0].linear.weight.dtype is torch.float32 + assert model.blocks[0].norm.weight.dtype is torch.bfloat16 + + +def test_exact_beats_wildcard_regardless_of_order(): + model = _model() + for rules in ( + (Rule("blocks.0.norm", gather="default"), Rule("blocks.*.norm", master="fp32", gather="fp32")), + (Rule("blocks.*.norm", master="fp32", gather="fp32"), Rule("blocks.0.norm", gather="default")), + ): + compiled = compile_precision_plan(model, PrecisionSpec(rules=rules), default_dtype=torch.bfloat16) + assert compiled.subshard_groups[0].modules == [model.blocks[1].norm] + + +def test_same_node_same_axis_tie_rejected(): + spec = PrecisionSpec( + rules=( + Rule("blocks.*.norm", gather="fp32"), + Rule("*.0.norm", gather="fp16"), + ) + ) + with pytest.raises(ValueError, match="tie"): + compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + + def test_buffer_only_module_master_cast(): model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32"),)) + spec = PrecisionSpec(rules=(Rule("blocks.*.rope", master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} assert compiled.subshard_groups == [] @@ -82,31 +117,19 @@ def test_buffer_only_module_master_cast(): def test_buffer_only_module_gather_is_skipped(): # Buffers are never gathered; a gather pin on a paramless module lowers to nothing. - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule("blocks.*.rope", gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.subshard_groups == [] def test_gather_matching_default_is_inline(): - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) + spec = PrecisionSpec(rules=(Rule("blocks.*.norm", gather="default"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.subshard_groups == [] assert compiled.master_casts == [] -def test_last_rule_wins_per_axis(): - model = _model() - spec = PrecisionSpec( - rules=( - Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"), - Rule(ModuleSel(fqn="blocks.0.norm"), master="default", gather="default"), - ) - ) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert compiled.subshard_groups[0].modules == [model.blocks[1].norm] - - def test_unmatched_rule_rejected(): - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) + spec = PrecisionSpec(rules=(Rule("no.such.module", master="fp32"),)) with pytest.raises(ValueError, match="matched no module"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) From dd9fe7b23b98e2a3cb5e8df6ff35b4d84892fe9a Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 21:07:34 -0700 Subject: [PATCH 18/50] revert(fsdp): restore ModuleSel selectors with last-match-wins folding --- miles/backends/fsdp_utils/precision.py | 101 +++++++----------- .../fsdp_utils/test_precision_plan.py | 67 ++++-------- 2 files changed, 63 insertions(+), 105 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 394c9143..6e9aa042 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -1,16 +1,11 @@ """Fine-grained weight-precision control for FSDP2, at module granularity. A family declares per-module dtype intent as PrecisionSpec rules on its -TrainPipelineConfig. A rule annotates the module-FQN tree nodes its ``path`` -matches (segment-wise glob, ``*`` never crosses dots) with one or both axes: +TrainPipelineConfig; each rule pins one or both axes (last match wins per axis): - master: resident dtype of the module's params/buffers (optimizer precision) - gather: dtype the params are cast to for FSDP all-gather + forward -Rule order carries no meaning; per axis each module resolves upward from itself: - - nearest wins: the closest annotated node on the self -> root chain applies - - most specific wins: on one node, the pattern with more literal segments - applies; a full tie is a compile error -Compute dtype is not managed here: the trainer autocasts the DiT forward, -model-boundary input dtypes are family policy applied by +The weight spec does not manage compute dtype: the trainer autocasts the DiT +forward, model-boundary input dtypes are family policy applied by ``apply_input_dtype_policy`` below, and op-level exceptions belong to the monkey-patch registry. @@ -18,9 +13,9 @@ PrecisionSpec rules | - (1) annotate matched tree nodes; a rule matching nothing is an error + (1) match each rule to modules; a rule covers its modules' whole subtrees | - (2) per module, resolve (master, gather) leaf-upward as above + (2) per module, fold covering rules -> (master, gather) intent | | master axis gather axis | | @@ -69,11 +64,19 @@ def _resolve_axis(axis: str | None, default_dtype: torch.dtype) -> torch.dtype | # --------------------------------------------------------------------------- +@dataclass(frozen=True) +class ModuleSel: + """Matches modules (and their subtrees) by fnmatch-ing FQN and/or exact class name.""" + + fqn: str | None = None + cls: str | None = None + + @dataclass(frozen=True) class Rule: - """path: segment-wise glob over module FQNs; axes take a dtype name, "default", or None (untouched).""" + """Axes take a dtype name ("fp32"/"bf16"/"fp16"), "default" (the run's default dtype), or None (untouched).""" - path: str + select: ModuleSel master: str | None = None gather: str | None = None @@ -117,36 +120,22 @@ def apply_master_casts(self) -> None: def _validate_rule(rule: Rule) -> None: if rule.master is None and rule.gather is None: raise ValueError(f"precision rule sets no dtype axis: {rule}") + if rule.select.fqn is None and rule.select.cls is None: + raise ValueError(f"precision rule has an empty ModuleSel: {rule}") for axis in (rule.master, rule.gather): if axis is not None and axis != "default" and axis not in _DTYPES: raise ValueError(f"precision rule has unknown dtype {axis!r}: {rule}") -def _path_matches(path: str, fqn: str) -> bool: - if path == "" or fqn == "": - return path == fqn - pattern_segments = path.split(".") - fqn_segments = fqn.split(".") - return len(pattern_segments) == len(fqn_segments) and all( - fnmatch(seg, pat) for seg, pat in zip(fqn_segments, pattern_segments, strict=True) - ) - - -def _specificity(path: str) -> int: - return sum(1 for seg in path.split(".") if not any(c in seg for c in "*?[")) - - -def _resolve_node_axis(rules_at_node: list[Rule], axis_name: str, mod_fqn: str) -> str | None: - """Most-specific-wins on one node; a full tie between distinct values is a compile error.""" - setters = [rule for rule in rules_at_node if getattr(rule, axis_name) is not None] - if not setters: - return None - best = max(_specificity(rule.path) for rule in setters) - winners = [rule for rule in setters if _specificity(rule.path) == best] - values = {getattr(rule, axis_name) for rule in winners} - if len(values) > 1: - raise ValueError(f"precision rules tie on {mod_fqn}.{axis_name}: {winners}") - return values.pop() +def _matched_module_prefixes(sel: ModuleSel, model: torch.nn.Module) -> list[str]: + prefixes = [] + for mod_fqn, module in model.named_modules(): + if sel.cls is not None and type(module).__name__ != sel.cls: + continue + if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): + continue + prefixes.append(mod_fqn) + return prefixes def compile_precision_plan( @@ -159,33 +148,25 @@ def compile_precision_plan( for rule in spec.rules: _validate_rule(rule) - # (1) Annotate matched tree nodes; a rule that matches nothing is almost certainly a typo. - all_fqns = [fqn for fqn, _ in model.named_modules()] - annotations: dict[str, list[Rule]] = {} - for rule in spec.rules: - matched = [fqn for fqn in all_fqns if _path_matches(rule.path, fqn)] - if not matched: + # (1) Match each rule to modules once; a rule that matches nothing is almost certainly a typo. + rule_prefixes = [_matched_module_prefixes(rule.select, model) for rule in spec.rules] + for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): + if not prefixes: raise ValueError(f"precision rule matched no module: {rule}") - for fqn in matched: - annotations.setdefault(fqn, []).append(rule) - - def _resolve(mod_fqn: str, axis_name: str) -> str | None: - node = mod_fqn - while True: - if node in annotations: - value = _resolve_node_axis(annotations[node], axis_name, node) - if value is not None: - return value - if node == "": - return None - node = node.rsplit(".", 1)[0] if "." in node else "" - - # (2)-(4) Resolve each module leaf-upward, lower master to casts and gather to sub-shard groups. + + def _covers(prefixes: list[str], mod_fqn: str) -> bool: + return any(p == "" or mod_fqn == p or mod_fqn.startswith(f"{p}.") for p in prefixes) + + # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard groups. master_casts: list[MasterCast] = [] groups: dict[torch.dtype, SubShardGroup] = {} for mod_fqn, module in model.named_modules(): - master = _resolve(mod_fqn, "master") - gather = _resolve(mod_fqn, "gather") + master = gather = None + for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): + if not _covers(prefixes, mod_fqn): + continue + master = rule.master if rule.master is not None else master + gather = rule.gather if rule.gather is not None else gather if master is None and gather is None: continue diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 8a48e0f6..883d889c 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn -from miles.backends.fsdp_utils.precision import PrecisionSpec, Rule, compile_precision_plan +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision_plan class Rope(nn.Module): @@ -39,9 +39,9 @@ def test_empty_spec_compiles_to_nothing(): assert compiled.subshard_groups == [] -def test_wildcard_rule_lowers_norms_to_one_subshard_group(): +def test_cls_rule_lowers_norms_to_one_subshard_group(): model = _model() - spec = PrecisionSpec(rules=(Rule("blocks.*.norm", master="fp32", gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert len(compiled.subshard_groups) == 1 group = compiled.subshard_groups[0] @@ -55,15 +55,15 @@ def test_wildcard_rule_lowers_norms_to_one_subshard_group(): def test_gather_without_master_is_allowed(): # Nested policy casts at all-gather, so gather may diverge from master. - spec = PrecisionSpec(rules=(Rule("blocks.*.norm", gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert compiled.subshard_groups[0].param_dtype is torch.float32 -def test_ancestor_rule_covers_subtree(): +def test_master_rule_covers_matched_subtree(): model = _model() - spec = PrecisionSpec(rules=(Rule("blocks.1", master="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert all(cast.fqn.startswith("blocks.1.") for cast in compiled.master_casts) compiled.apply_master_casts() @@ -72,44 +72,9 @@ def test_ancestor_rule_covers_subtree(): assert model.blocks[0].linear.weight.dtype is torch.bfloat16 -def test_nearest_annotation_wins(): - model = _model() - spec = PrecisionSpec( - rules=( - Rule("blocks.0.norm", master="default"), - Rule("blocks.0", master="fp32"), - ) - ) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - compiled.apply_master_casts() - assert model.blocks[0].linear.weight.dtype is torch.float32 - assert model.blocks[0].norm.weight.dtype is torch.bfloat16 - - -def test_exact_beats_wildcard_regardless_of_order(): - model = _model() - for rules in ( - (Rule("blocks.0.norm", gather="default"), Rule("blocks.*.norm", master="fp32", gather="fp32")), - (Rule("blocks.*.norm", master="fp32", gather="fp32"), Rule("blocks.0.norm", gather="default")), - ): - compiled = compile_precision_plan(model, PrecisionSpec(rules=rules), default_dtype=torch.bfloat16) - assert compiled.subshard_groups[0].modules == [model.blocks[1].norm] - - -def test_same_node_same_axis_tie_rejected(): - spec = PrecisionSpec( - rules=( - Rule("blocks.*.norm", gather="fp32"), - Rule("*.0.norm", gather="fp16"), - ) - ) - with pytest.raises(ValueError, match="tie"): - compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) - - def test_buffer_only_module_master_cast(): model = _model() - spec = PrecisionSpec(rules=(Rule("blocks.*.rope", master="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} assert compiled.subshard_groups == [] @@ -117,19 +82,31 @@ def test_buffer_only_module_master_cast(): def test_buffer_only_module_gather_is_skipped(): # Buffers are never gathered; a gather pin on a paramless module lowers to nothing. - spec = PrecisionSpec(rules=(Rule("blocks.*.rope", gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.subshard_groups == [] def test_gather_matching_default_is_inline(): - spec = PrecisionSpec(rules=(Rule("blocks.*.norm", gather="default"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.subshard_groups == [] assert compiled.master_casts == [] +def test_last_rule_wins_per_axis(): + model = _model() + spec = PrecisionSpec( + rules=( + Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.norm"), master="default", gather="default"), + ) + ) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert compiled.subshard_groups[0].modules == [model.blocks[1].norm] + + def test_unmatched_rule_rejected(): - spec = PrecisionSpec(rules=(Rule("no.such.module", master="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) with pytest.raises(ValueError, match="matched no module"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) From 4bf099f4c2c35d4cf08dca318427bee5b7fac395 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 21:09:58 -0700 Subject: [PATCH 19/50] feat(fsdp): wrap sub-shard groups in bottom-up topological order --- miles/backends/fsdp_utils/precision.py | 46 +++++++++++++----- .../fsdp_utils/test_precision_plan.py | 47 +++++++++++++++---- 2 files changed, 72 insertions(+), 21 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 6e9aa042..1e2ee208 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -26,9 +26,9 @@ SubShardGroup per dtype (buffer-only and paramless modules have nothing to gather: skipped) | - apply_fsdp2: fully_shard(group.modules, param_dtype=group dtype), - nested before the block/root wrap; one extra all-gather per group per - step (reshard_after_forward=False, so backward re-uses the forward gather) + apply_fsdp2: fully_shard(group.modules, param_dtype=group dtype), deepest + group first then blocks then root (FSDP2 nests bottom-up); one extra + all-gather per group per step (reshard_after_forward=False) Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and FSDP2 requires uniform master dtype among trainable params per unit), so @@ -157,9 +157,9 @@ def compile_precision_plan( def _covers(prefixes: list[str], mod_fqn: str) -> bool: return any(p == "" or mod_fqn == p or mod_fqn.startswith(f"{p}.") for p in prefixes) - # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard groups. + # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard requests. master_casts: list[MasterCast] = [] - groups: dict[torch.dtype, SubShardGroup] = {} + pending: dict[torch.dtype, list[tuple[str, torch.nn.Module, list[str]]]] = {} for mod_fqn, module in model.named_modules(): master = gather = None for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): @@ -181,16 +181,40 @@ def _covers(prefixes: list[str], mod_fqn: str) -> bool: gather_dtype = _resolve_axis(gather, default_dtype) if gather_dtype is not None and gather_dtype != default_dtype: - float_params = [name for name, param in params if param.is_floating_point()] + float_params = [f"{prefix}{name}" for name, param in params if param.is_floating_point()] if not float_params: continue if mod_fqn == "": raise ValueError("cannot sub-wrap the root module for a gather override") - group = groups.setdefault(gather_dtype, SubShardGroup(param_dtype=gather_dtype)) - group.modules.append(module) - group.param_fqns.extend(f"{prefix}{name}" for name in float_params) - - return CompiledPrecision(master_casts=master_casts, subshard_groups=list(groups.values())) + pending.setdefault(gather_dtype, []).append((mod_fqn, module, float_params)) + + # (5) Bottom-up wrap order: same-dtype descendants coalesce into their ancestor entry + # (named_modules is parent-first), groups sort deepest-first, and cross-dtype nesting + # that cannot wrap child-before-parent is rejected. + built: list[tuple[int, list[str], SubShardGroup]] = [] + for dtype, entries in pending.items(): + kept: list[tuple[str, torch.nn.Module, list[str]]] = [] + for mod_fqn, module, float_params in entries: + ancestor = next((entry for entry in kept if mod_fqn.startswith(f"{entry[0]}.")), None) + if ancestor is not None: + ancestor[2].extend(float_params) + continue + kept.append((mod_fqn, module, float_params)) + depth = max(mod_fqn.count(".") for mod_fqn, _, _ in kept) + group = SubShardGroup( + param_dtype=dtype, + modules=[module for _, module, _ in kept], + param_fqns=[fqn for _, _, fqns in kept for fqn in fqns], + ) + built.append((depth, [mod_fqn for mod_fqn, _, _ in kept], group)) + built.sort(key=lambda item: -item[0]) + for i, (_, fqns_i, _) in enumerate(built): + for _, fqns_j, _ in built[i + 1 :]: + for fqn_j in fqns_j: + if any(fqn_j.startswith(f"{fqn_i}.") for fqn_i in fqns_i): + raise ValueError(f"sub-shard nesting cannot wrap bottom-up: {fqn_j} sits inside an earlier group") + + return CompiledPrecision(master_casts=master_casts, subshard_groups=[group for _, _, group in built]) # --------------------------------------------------------------------------- diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 883d889c..b9cea9f2 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -15,11 +15,19 @@ def __init__(self): self.register_buffer("freqs", torch.zeros(4)) +class Attn(nn.Module): + def __init__(self): + super().__init__() + self.norm_q = nn.LayerNorm(8) + + class Block(nn.Module): def __init__(self): super().__init__() + self.scale = nn.Parameter(torch.zeros(8)) self.linear = nn.Linear(8, 8) self.norm = nn.LayerNorm(8) + self.attn = Attn() self.rope = Rope() @@ -33,32 +41,35 @@ def _model(dtype=torch.bfloat16): return Tiny().to(dtype) +def _norms(model): + return [model.blocks[0].norm, model.blocks[0].attn.norm_q, model.blocks[1].norm, model.blocks[1].attn.norm_q] + + def test_empty_spec_compiles_to_nothing(): compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert compiled.subshard_groups == [] -def test_cls_rule_lowers_norms_to_one_subshard_group(): +def test_fqn_wildcard_selects_norms_across_depths(): model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert len(compiled.subshard_groups) == 1 group = compiled.subshard_groups[0] assert group.param_dtype is torch.float32 - assert group.modules == [model.blocks[0].norm, model.blocks[1].norm] - assert set(group.param_fqns) == {f"blocks.{i}.norm.{n}" for i in range(2) for n in ("weight", "bias")} + assert group.modules == _norms(model) compiled.apply_master_casts() - assert model.blocks[0].norm.weight.dtype is torch.float32 + assert model.blocks[0].attn.norm_q.weight.dtype is torch.float32 assert model.blocks[0].linear.weight.dtype is torch.bfloat16 -def test_gather_without_master_is_allowed(): - # Nested policy casts at all-gather, so gather may diverge from master. +def test_cls_rule_selects_by_class(): + model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] - assert compiled.subshard_groups[0].param_dtype is torch.float32 + assert compiled.subshard_groups[0].modules == _norms(model) def test_master_rule_covers_matched_subtree(): @@ -103,7 +114,23 @@ def test_last_rule_wins_per_axis(): ) ) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert compiled.subshard_groups[0].modules == [model.blocks[1].norm] + assert compiled.subshard_groups[0].modules == _norms(model)[1:] + + +def test_nested_gather_groups_wrap_bottom_up(): + model = _model() + spec = PrecisionSpec( + rules=( + Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), + Rule(ModuleSel(fqn="blocks.0.norm"), gather="fp32"), + ) + ) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + # Deeper fp32 norm group wraps first; same-dtype children coalesce into blocks.0. + assert [group.param_dtype for group in compiled.subshard_groups] == [torch.float32, torch.float16] + assert compiled.subshard_groups[0].modules == [model.blocks[0].norm] + assert compiled.subshard_groups[1].modules == [model.blocks[0]] + assert "blocks.0.linear.weight" in compiled.subshard_groups[1].param_fqns def test_unmatched_rule_rejected(): From cd5911d04605879018df58a333a333200a0a5f11 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 21:34:56 -0700 Subject: [PATCH 20/50] refactor(fsdp): make each selected module its own FSDP2 wrap unit --- miles/backends/fsdp_utils/actor.py | 23 +- miles/backends/fsdp_utils/precision.py | 206 +++++++++--------- .../fsdp_utils/_precision_wrap_worker.py | 128 +++++++++++ .../fsdp_utils/test_precision_plan.py | 100 ++++----- .../fsdp_utils/test_precision_wrap.py | 37 ++++ 5 files changed, 327 insertions(+), 167 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/_precision_wrap_worker.py create mode 100644 tests/fast/backends/fsdp_utils/test_precision_wrap.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 13999208..d3aa1ea7 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -154,7 +154,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), - subshard_groups=compiled_precision.subshard_groups, + precision_wrap_units=compiled_precision.wrap_units, ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -624,7 +624,7 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, subshard_groups=None): +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, precision_wrap_units=None): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard offload_policy = CPUOffloadPolicy() if cpu_offload else None @@ -638,7 +638,7 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules reduce_dtype = resolve_dtype(args.fsdp_reduce_dtype) logger.info( f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, " - f"reduce_dtype={reduce_dtype}, sub-shard groups={len(subshard_groups) if subshard_groups else 0}" + f"reduce_dtype={reduce_dtype}, precision wrap units={len(precision_wrap_units) if precision_wrap_units else 0}" ) def _fsdp_kwargs(policy_param_dtype): @@ -653,18 +653,17 @@ def _fsdp_kwargs(policy_param_dtype): "mesh": mesh, } - # Sub-shard groups wrap first so the block/root units exclude their params. - # Groups are tiny by design: ZeRO-2 style (no backward re-gather) costs a few MB. - subshard_modules = set() - for group in subshard_groups or (): - fully_shard(group.modules, reshard_after_forward=False, **_fsdp_kwargs(group.param_dtype)) - subshard_modules.update(group.modules) + # Precision units come pre-sorted deepest-first, so every wrap excludes the ones already wrapped. + precision_modules = set() + for unit in precision_wrap_units or (): + fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype)) + precision_modules.add(unit.module) fsdp_kwargs = _fsdp_kwargs(param_dtype) for module in modules: - if module in subshard_modules: - raise ValueError(f"{type(module).__name__} is both an FSDP block unit and a precision sub-shard target") - fully_shard(module, **fsdp_kwargs) + # A block that is also a precision unit keeps its pinned policy; wrapping twice is an error. + if module not in precision_modules: + fully_shard(module, **fsdp_kwargs) fully_shard(model, **fsdp_kwargs) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 1e2ee208..859c89ae 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -1,11 +1,13 @@ """Fine-grained weight-precision control for FSDP2, at module granularity. -A family declares per-module dtype intent as PrecisionSpec rules on its -TrainPipelineConfig; each rule pins one or both axes (last match wins per axis): +A family declares dtype intent as PrecisionSpec rules on its TrainPipelineConfig. +A rule selects modules by FQN glob and/or class-name glob (both narrows to the +intersection) and pins one or both axes; a rule covers its modules' subtrees, and +where rules overlap the later one wins per axis: - master: resident dtype of the module's params/buffers (optimizer precision) - gather: dtype the params are cast to for FSDP all-gather + forward -The weight spec does not manage compute dtype: the trainer autocasts the DiT -forward, model-boundary input dtypes are family policy applied by +Compute dtype is not managed here: the trainer autocasts the DiT forward, +model-boundary input dtypes are family policy applied by ``apply_input_dtype_policy`` below, and op-level exceptions belong to the monkey-patch registry. @@ -13,32 +15,34 @@ PrecisionSpec rules | - (1) match each rule to modules; a rule covers its modules' whole subtrees + (1) match each rule to modules once | (2) per module, fold covering rules -> (master, gather) intent | | master axis gather axis | | - (3) != loaded dtype? (4) != default dtype? --no--> inline: the block - -> MasterCast of the | policy already is - module's own float v the default - tensors, at load time sub-wrap: merge same-dtype modules into one - SubShardGroup per dtype (buffer-only and - paramless modules have nothing to gather: skipped) - | - apply_fsdp2: fully_shard(group.modules, param_dtype=group dtype), deepest - group first then blocks then root (FSDP2 nests bottom-up); one extra - all-gather per group per step (reshard_after_forward=False) - -Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, -and FSDP2 requires uniform master dtype among trainable params per unit), so + (3) != loaded dtype? (4) != the dtype the enclosing wrap unit already + -> MasterCast of the provides? -> the module becomes its own wrap + module's own float unit at that dtype (paramless modules have + tensors, at load time nothing to gather and are skipped) + | | + v v + apply_master_casts() apply_fsdp2: fully_shard(unit.module, + before FSDP wrapping param_dtype=unit dtype), deepest unit first, then + blocks, then root. FSDP2 nests child-before-parent, + so a unit is always excluded from its enclosing + unit — that is how gather="default" carves a module + back out of a non-default ancestor. + +Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and +FSDP2 requires uniform master dtype among trainable params per unit), so finer-grained selectors are deliberately not offered. """ from __future__ import annotations import logging -from dataclasses import dataclass, field +from dataclasses import dataclass from fnmatch import fnmatch import torch @@ -66,7 +70,7 @@ def _resolve_axis(axis: str | None, default_dtype: torch.dtype) -> torch.dtype | @dataclass(frozen=True) class ModuleSel: - """Matches modules (and their subtrees) by fnmatch-ing FQN and/or exact class name.""" + """Module selector; fqn and cls are globs over the module FQN and class name.""" fqn: str | None = None cls: str | None = None @@ -87,7 +91,7 @@ class PrecisionSpec: # --------------------------------------------------------------------------- -# Compiler: per-module plan -> FSDP2 lowering (master casts + sub-shard groups) +# Compiler: spec -> FSDP2 lowering (master casts + per-module wrap units) # --------------------------------------------------------------------------- @@ -98,19 +102,19 @@ class MasterCast: dtype: torch.dtype -@dataclass -class SubShardGroup: - """Modules to wrap as one nested fully_shard unit with param_dtype=gather.""" +@dataclass(frozen=True) +class WrapUnit: + """A module to fully_shard on its own with param_dtype=gather.""" + fqn: str + module: torch.nn.Module param_dtype: torch.dtype - modules: list[torch.nn.Module] = field(default_factory=list) - param_fqns: list[str] = field(default_factory=list) @dataclass class CompiledPrecision: master_casts: list[MasterCast] - subshard_groups: list[SubShardGroup] + wrap_units: list[WrapUnit] def apply_master_casts(self) -> None: for cast in self.master_casts: @@ -127,15 +131,39 @@ def _validate_rule(rule: Rule) -> None: raise ValueError(f"precision rule has unknown dtype {axis!r}: {rule}") -def _matched_module_prefixes(sel: ModuleSel, model: torch.nn.Module) -> list[str]: - prefixes = [] - for mod_fqn, module in model.named_modules(): - if sel.cls is not None and type(module).__name__ != sel.cls: - continue - if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): +def _selects(sel: ModuleSel, mod_fqn: str, module: torch.nn.Module) -> bool: + if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): + return False + return sel.cls is None or fnmatch(type(module).__name__, sel.cls) + + +def _parent_fqn(mod_fqn: str) -> str: + """The root module's FQN is "", and it is its own parent.""" + return mod_fqn.rsplit(".", 1)[0] if "." in mod_fqn else "" + + +def _self_and_ancestors(mod_fqn: str) -> list[str]: + fqns = [] + while mod_fqn: + fqns.append(mod_fqn) + mod_fqn = _parent_fqn(mod_fqn) + return fqns + [""] + + +def _fold_covering_rules( + rules: tuple[Rule, ...], + matched_fqns: list[set[str]], + mod_fqn: str, +) -> tuple[str | None, str | None]: + """Later rules override earlier ones per axis; a rule covers its matched modules' subtrees.""" + covering = _self_and_ancestors(mod_fqn) + master = gather = None + for rule, fqns in zip(rules, matched_fqns, strict=True): + if fqns.isdisjoint(covering): continue - prefixes.append(mod_fqn) - return prefixes + master = rule.master if rule.master is not None else master + gather = rule.gather if rule.gather is not None else gather + return master, gather def compile_precision_plan( @@ -148,73 +176,55 @@ def compile_precision_plan( for rule in spec.rules: _validate_rule(rule) - # (1) Match each rule to modules once; a rule that matches nothing is almost certainly a typo. - rule_prefixes = [_matched_module_prefixes(rule.select, model) for rule in spec.rules] - for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): - if not prefixes: + # (1) A rule matching nothing is almost certainly a typo'd pattern or class name. + matched_fqns = [] + for rule in spec.rules: + fqns = {fqn for fqn, module in model.named_modules() if _selects(rule.select, fqn, module)} + if not fqns: raise ValueError(f"precision rule matched no module: {rule}") + matched_fqns.append(fqns) - def _covers(prefixes: list[str], mod_fqn: str) -> bool: - return any(p == "" or mod_fqn == p or mod_fqn.startswith(f"{p}.") for p in prefixes) - - # (2)-(4) Fold rules per module, lower master to casts and gather to sub-shard requests. + # (2)-(4) named_modules is parent-first, so the enclosing unit's dtype is already known here. master_casts: list[MasterCast] = [] - pending: dict[torch.dtype, list[tuple[str, torch.nn.Module, list[str]]]] = {} + wrap_units: list[WrapUnit] = [] + unit_dtypes: dict[str, torch.dtype] = {"": default_dtype} for mod_fqn, module in model.named_modules(): - master = gather = None - for rule, prefixes in zip(spec.rules, rule_prefixes, strict=True): - if not _covers(prefixes, mod_fqn): - continue - master = rule.master if rule.master is not None else master - gather = rule.gather if rule.gather is not None else gather - if master is None and gather is None: - continue + enclosing_dtype = unit_dtypes[_parent_fqn(mod_fqn)] + unit_dtypes[mod_fqn] = enclosing_dtype + master, gather = _fold_covering_rules(spec.rules, matched_fqns, mod_fqn) - prefix = f"{mod_fqn}." if mod_fqn else "" - params = list(module.named_parameters(recurse=False)) - tensors = params + list(module.named_buffers(recurse=False)) master_dtype = _resolve_axis(master, default_dtype) if master_dtype is not None: - for name, tensor in tensors: + prefix = f"{mod_fqn}." if mod_fqn else "" + own = list(module.named_parameters(recurse=False)) + list(module.named_buffers(recurse=False)) + for name, tensor in own: if tensor.is_floating_point() and tensor.dtype != master_dtype: master_casts.append(MasterCast(f"{prefix}{name}", tensor, master_dtype)) gather_dtype = _resolve_axis(gather, default_dtype) - if gather_dtype is not None and gather_dtype != default_dtype: - float_params = [f"{prefix}{name}" for name, param in params if param.is_floating_point()] - if not float_params: - continue - if mod_fqn == "": - raise ValueError("cannot sub-wrap the root module for a gather override") - pending.setdefault(gather_dtype, []).append((mod_fqn, module, float_params)) - - # (5) Bottom-up wrap order: same-dtype descendants coalesce into their ancestor entry - # (named_modules is parent-first), groups sort deepest-first, and cross-dtype nesting - # that cannot wrap child-before-parent is rejected. - built: list[tuple[int, list[str], SubShardGroup]] = [] - for dtype, entries in pending.items(): - kept: list[tuple[str, torch.nn.Module, list[str]]] = [] - for mod_fqn, module, float_params in entries: - ancestor = next((entry for entry in kept if mod_fqn.startswith(f"{entry[0]}.")), None) - if ancestor is not None: - ancestor[2].extend(float_params) - continue - kept.append((mod_fqn, module, float_params)) - depth = max(mod_fqn.count(".") for mod_fqn, _, _ in kept) - group = SubShardGroup( - param_dtype=dtype, - modules=[module for _, module, _ in kept], - param_fqns=[fqn for _, _, fqns in kept for fqn in fqns], - ) - built.append((depth, [mod_fqn for mod_fqn, _, _ in kept], group)) - built.sort(key=lambda item: -item[0]) - for i, (_, fqns_i, _) in enumerate(built): - for _, fqns_j, _ in built[i + 1 :]: - for fqn_j in fqns_j: - if any(fqn_j.startswith(f"{fqn_i}.") for fqn_i in fqns_i): - raise ValueError(f"sub-shard nesting cannot wrap bottom-up: {fqn_j} sits inside an earlier group") - - return CompiledPrecision(master_casts=master_casts, subshard_groups=[group for _, _, group in built]) + if gather_dtype is None or gather_dtype == enclosing_dtype: + continue + if not any(param.is_floating_point() for param in module.parameters()): + continue + if mod_fqn == "": + raise ValueError("cannot wrap the root module for a gather override") + wrap_units.append(WrapUnit(mod_fqn, module, gather_dtype)) + unit_dtypes[mod_fqn] = gather_dtype + + # (5) FSDP2 nests child-before-parent, so hand back the deepest units first. + wrap_units.sort(key=lambda unit: -unit.fqn.count(".")) + return CompiledPrecision(master_casts=master_casts, wrap_units=wrap_units) + + +def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: + logger.info( + f"precision[{component}]: default gather dtype {default_dtype}, " + f"{len(compiled.master_casts)} master casts, {len(compiled.wrap_units)} extra wrap units" + ) + for cast in compiled.master_casts: + logger.info(f"precision[{component}]: master {cast.fqn} -> {cast.dtype}") + for unit in compiled.wrap_units: + logger.info(f"precision[{component}]: wrap {unit.fqn} @ {unit.param_dtype}") # --------------------------------------------------------------------------- @@ -254,17 +264,3 @@ def _cast(value, dtype: torch.dtype | None): None if cond is None else {key: _cast(value, cond_dtype) for key, value in cond.items()} for cond in conds ) return _cast(latents, latents_dtype), _cast(timesteps, timestep_dtype), out_conds - - -def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: - logger.info( - f"precision[{component}]: default gather dtype {default_dtype}, " - f"{len(compiled.master_casts)} master casts, {len(compiled.subshard_groups)} sub-shard groups" - ) - for cast in compiled.master_casts: - logger.info(f"precision[{component}]: master {cast.fqn} -> {cast.dtype}") - for group in compiled.subshard_groups: - logger.info( - f"precision[{component}]: sub-shard @ {group.param_dtype}: " - f"{len(group.modules)} modules, params {group.param_fqns}" - ) diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py new file mode 100644 index 00000000..00150817 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -0,0 +1,128 @@ +"""Gloo worker asserting the compiled precision plan really wraps under FSDP2 (2 ranks). + +Modules cast explicitly in forward because CPU kernels reject mixed dtypes; what +matters here is the wrap nesting and the param dtype each module sees at forward. +""" + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard + +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision_plan + +DEFAULT_DTYPE = torch.bfloat16 + + +class Norm(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(8)) + + def forward(self, x): + return x * self.weight.to(x.dtype) + + +class Proj(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.eye(8)) + + def forward(self, x): + return x @ self.weight.to(x.dtype) + + +class Attn(nn.Module): + def __init__(self): + super().__init__() + self.norm_q = Norm() + self.proj = Proj() + + def forward(self, x): + return self.proj(self.norm_q(x)) + + +class Block(nn.Module): + def __init__(self): + super().__init__() + self.norm = Norm() + self.attn = Attn() + + def forward(self, x): + return self.attn(self.norm(x)) + + +class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([Block(), Block()]) + + def forward(self, x): + for block in self.blocks: + x = block(x) + return x + + +SPEC = PrecisionSpec( + rules=( + Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), + Rule(ModuleSel(fqn="blocks.0.attn"), gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.attn.norm_q"), gather="default"), + Rule(ModuleSel(cls="Norm", fqn="blocks.1*"), gather="fp32"), + ) +) +EXPECTED_GATHER = { + "blocks.0.norm": torch.float16, # inherits the fp16 block unit + "blocks.0.attn.norm_q": DEFAULT_DTYPE, # carved back out of two non-default ancestors + "blocks.0.attn.proj": torch.float32, # inherits the fp32 attn unit + "blocks.1.norm": torch.float32, # cls + fqn rule + "blocks.1.attn.norm_q": torch.float32, + "blocks.1.attn.proj": DEFAULT_DTYPE, # untouched by any rule +} + + +def main() -> None: + dist.init_process_group("gloo") + mesh = init_device_mesh("cpu", (dist.get_world_size(),)) + model = Tiny().to(torch.float32) # fp32 master + + compiled = compile_precision_plan(model, SPEC, default_dtype=DEFAULT_DTYPE) + compiled.apply_master_casts() + + def fsdp_kwargs(param_dtype): + policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) + return {"mp_policy": policy, "mesh": mesh} + + precision_modules = set() + for unit in compiled.wrap_units: + fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype)) + precision_modules.add(unit.module) + for block in model.blocks: + if block not in precision_modules: + fully_shard(block, **fsdp_kwargs(DEFAULT_DTYPE)) + fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE)) + + seen: dict[str, torch.dtype] = {} + + def record(module, _args, fqn): + seen.setdefault(fqn, next(module.parameters(recurse=False)).dtype) + return None + + for fqn, module in model.named_modules(): + if list(module.parameters(recurse=False)): + module.register_forward_pre_hook(lambda module, args, fqn=fqn: record(module, args, fqn)) + + model(torch.randn(2, 8, dtype=DEFAULT_DTYPE)).float().sum().backward() + + for fqn, want in EXPECTED_GATHER.items(): + if seen.get(fqn) != want: + raise AssertionError(f"{fqn} gathered as {seen.get(fqn)}, expected {want}") + weight = model.blocks[0].attn.norm_q.weight + if weight.dtype != torch.float32 or weight.grad.dtype != torch.float32: + raise AssertionError(f"master/grad left fp32: {weight.dtype}/{weight.grad.dtype}") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index b9cea9f2..26d826da 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -24,7 +24,6 @@ def __init__(self): class Block(nn.Module): def __init__(self): super().__init__() - self.scale = nn.Parameter(torch.zeros(8)) self.linear = nn.Linear(8, 8) self.norm = nn.LayerNorm(8) self.attn = Attn() @@ -41,96 +40,97 @@ def _model(dtype=torch.bfloat16): return Tiny().to(dtype) -def _norms(model): - return [model.blocks[0].norm, model.blocks[0].attn.norm_q, model.blocks[1].norm, model.blocks[1].attn.norm_q] +def _units(compiled): + return [(unit.fqn, unit.param_dtype) for unit in compiled.wrap_units] + + +NORM_FQNS = {f"blocks.{i}{suffix}" for i in range(2) for suffix in (".norm", ".attn.norm_q")} def test_empty_spec_compiles_to_nothing(): compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.master_casts == [] - assert compiled.subshard_groups == [] + assert compiled.wrap_units == [] -def test_fqn_wildcard_selects_norms_across_depths(): +def test_fqn_glob_selects_norms_across_depths(): model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert len(compiled.subshard_groups) == 1 - group = compiled.subshard_groups[0] - assert group.param_dtype is torch.float32 - assert group.modules == _norms(model) + assert dict(_units(compiled)) == dict.fromkeys(NORM_FQNS, torch.float32) compiled.apply_master_casts() assert model.blocks[0].attn.norm_q.weight.dtype is torch.float32 assert model.blocks[0].linear.weight.dtype is torch.bfloat16 -def test_cls_rule_selects_by_class(): - model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="fp32"),)) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) +def test_cls_glob_selects_by_class(): + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="*LayerNorm"), gather="fp32"),)) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] - assert compiled.subshard_groups[0].modules == _norms(model) + assert {fqn for fqn, _ in _units(compiled)} == NORM_FQNS def test_master_rule_covers_matched_subtree(): model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert all(cast.fqn.startswith("blocks.1.") for cast in compiled.master_casts) compiled.apply_master_casts() assert model.blocks[1].linear.weight.dtype is torch.float32 assert model.blocks[1].rope.freqs.dtype is torch.float32 assert model.blocks[0].linear.weight.dtype is torch.bfloat16 -def test_buffer_only_module_master_cast(): - model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32"),)) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} - assert compiled.subshard_groups == [] - - -def test_buffer_only_module_gather_is_skipped(): - # Buffers are never gathered; a gather pin on a paramless module lowers to nothing. - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), gather="fp32"),)) +def test_later_rule_overrides_earlier_selection(): + spec = PrecisionSpec( + rules=( + Rule(ModuleSel(cls="LayerNorm"), gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.*norm*"), gather="fp16"), + ) + ) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) - assert compiled.subshard_groups == [] + assert dict(_units(compiled)) == { + "blocks.0.norm": torch.float16, + "blocks.0.attn.norm_q": torch.float16, + "blocks.1.norm": torch.float32, + "blocks.1.attn.norm_q": torch.float32, + } -def test_gather_matching_default_is_inline(): - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="LayerNorm"), gather="default"),)) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) - assert compiled.subshard_groups == [] - assert compiled.master_casts == [] +def test_empty_module_sel_rejected(): + spec = PrecisionSpec(rules=(Rule(ModuleSel(), master="fp32"),)) + with pytest.raises(ValueError, match="empty ModuleSel"): + compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) -def test_last_rule_wins_per_axis(): - model = _model() +def test_every_node_of_a_nested_chain_wraps_bottom_up(): spec = PrecisionSpec( rules=( - Rule(ModuleSel(cls="LayerNorm"), master="fp32", gather="fp32"), - Rule(ModuleSel(fqn="blocks.0.norm"), master="default", gather="default"), + Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), + Rule(ModuleSel(fqn="blocks.0.attn"), gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.attn.norm_q"), gather="default"), ) ) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - assert compiled.subshard_groups[0].modules == _norms(model)[1:] + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert _units(compiled) == [ + ("blocks.0.attn.norm_q", torch.bfloat16), + ("blocks.0.attn", torch.float32), + ("blocks.0", torch.float16), + ] -def test_nested_gather_groups_wrap_bottom_up(): +def test_inherited_gather_needs_no_extra_unit(): + # attn owns no parameter but its subtree does, so it wraps; norm_q inherits the same dtype. + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Attn"), gather="fp32"),)) + compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + assert _units(compiled) == [("blocks.0.attn", torch.float32), ("blocks.1.attn", torch.float32)] + + +def test_buffer_only_module_casts_master_without_wrapping(): model = _model() - spec = PrecisionSpec( - rules=( - Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), - Rule(ModuleSel(fqn="blocks.0.norm"), gather="fp32"), - ) - ) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) - # Deeper fp32 norm group wraps first; same-dtype children coalesce into blocks.0. - assert [group.param_dtype for group in compiled.subshard_groups] == [torch.float32, torch.float16] - assert compiled.subshard_groups[0].modules == [model.blocks[0].norm] - assert compiled.subshard_groups[1].modules == [model.blocks[0]] - assert "blocks.0.linear.weight" in compiled.subshard_groups[1].param_fqns + assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} + assert compiled.wrap_units == [] def test_unmatched_rule_rejected(): diff --git a/tests/fast/backends/fsdp_utils/test_precision_wrap.py b/tests/fast/backends/fsdp_utils/test_precision_wrap.py new file mode 100644 index 00000000..c6b16989 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_precision_wrap.py @@ -0,0 +1,37 @@ +"""Nested precision wrap units under real FSDP2 on 2 gloo ranks; assertions live in the worker.""" + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu", labels=[]) + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +_WORKER = Path(__file__).with_name("_precision_wrap_worker.py") + + +def test_precision_wrap_units(): + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc_per_node=2", + str(_WORKER), + ], + check=True, + env=env, + timeout=300, + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) From bbdb334896fab117ad69ec0c9b2adceba98ec130 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 21:41:00 -0700 Subject: [PATCH 21/50] refactor(fsdp): merge precision units and block modules into one deepest-first wrap plan --- miles/backends/fsdp_utils/actor.py | 23 ++++++------ miles/backends/fsdp_utils/precision.py | 35 +++++++++++++------ .../fsdp_utils/_precision_wrap_worker.py | 12 +++---- .../fsdp_utils/test_precision_plan.py | 21 +++++++---- 4 files changed, 53 insertions(+), 38 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d3aa1ea7..4712cfe0 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -39,7 +39,13 @@ from .lr_scheduler import get_lr_scheduler from .metrics import new_metric_buffer from .parallel import create_fsdp_parallel_state -from .precision import apply_input_dtype_policy, compile_precision_plan, log_precision_summary, resolve_dtype +from .precision import ( + apply_input_dtype_policy, + build_wrap_plan, + compile_precision_plan, + log_precision_summary, + resolve_dtype, +) from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -653,18 +659,9 @@ def _fsdp_kwargs(policy_param_dtype): "mesh": mesh, } - # Precision units come pre-sorted deepest-first, so every wrap excludes the ones already wrapped. - precision_modules = set() - for unit in precision_wrap_units or (): - fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype)) - precision_modules.add(unit.module) - - fsdp_kwargs = _fsdp_kwargs(param_dtype) - for module in modules: - # A block that is also a precision unit keeps its pinned policy; wrapping twice is an error. - if module not in precision_modules: - fully_shard(module, **fsdp_kwargs) + for module, policy_dtype in build_wrap_plan(model, precision_wrap_units or [], modules, param_dtype): + fully_shard(module, **_fsdp_kwargs(policy_dtype)) - fully_shard(model, **fsdp_kwargs) + fully_shard(model, **_fsdp_kwargs(param_dtype)) return model diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 859c89ae..504032b0 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -24,15 +24,15 @@ (3) != loaded dtype? (4) != the dtype the enclosing wrap unit already -> MasterCast of the provides? -> the module becomes its own wrap module's own float unit at that dtype (paramless modules have - tensors, at load time nothing to gather and are skipped) - | | - v v - apply_master_casts() apply_fsdp2: fully_shard(unit.module, - before FSDP wrapping param_dtype=unit dtype), deepest unit first, then - blocks, then root. FSDP2 nests child-before-parent, - so a unit is always excluded from its enclosing - unit — that is how gather="default" carves a module - back out of a non-default ancestor. + tensors, at load time nothing to gather and are skipped), which makes + | the units a minimal cover of the tree + v | + apply_master_casts() v + before FSDP wrapping build_wrap_plan() merges the units with the block + modules into one deepest-first order, so FSDP2 + always nests child-before-parent — that is how + gather="default" carves a module back out of a + non-default ancestor. Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and FSDP2 requires uniform master dtype among trainable params per unit), so @@ -211,11 +211,24 @@ def compile_precision_plan( wrap_units.append(WrapUnit(mod_fqn, module, gather_dtype)) unit_dtypes[mod_fqn] = gather_dtype - # (5) FSDP2 nests child-before-parent, so hand back the deepest units first. - wrap_units.sort(key=lambda unit: -unit.fqn.count(".")) return CompiledPrecision(master_casts=master_casts, wrap_units=wrap_units) +def build_wrap_plan( + model: torch.nn.Module, + wrap_units: list[WrapUnit], + block_modules: list[torch.nn.Module], + default_dtype: torch.dtype, +) -> list[tuple[torch.nn.Module, torch.dtype]]: + """One wrap order for FSDP2, deepest module first: precision units pin their param_dtype, + remaining block modules take the default, and a module appearing in both keeps the pin.""" + plan: dict[torch.nn.Module, torch.dtype] = {unit.module: unit.param_dtype for unit in wrap_units} + for module in block_modules: + plan.setdefault(module, default_dtype) + depths = {module: mod_fqn.count(".") for mod_fqn, module in model.named_modules()} + return [(module, plan[module]) for module in sorted(plan, key=lambda module: -depths[module])] + + def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: logger.info( f"precision[{component}]: default gather dtype {default_dtype}, " diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index 00150817..068ae507 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -10,7 +10,7 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard -from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision_plan +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, build_wrap_plan, compile_precision_plan DEFAULT_DTYPE = torch.bfloat16 @@ -94,13 +94,9 @@ def fsdp_kwargs(param_dtype): policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) return {"mp_policy": policy, "mesh": mesh} - precision_modules = set() - for unit in compiled.wrap_units: - fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype)) - precision_modules.add(unit.module) - for block in model.blocks: - if block not in precision_modules: - fully_shard(block, **fsdp_kwargs(DEFAULT_DTYPE)) + plan = build_wrap_plan(model, compiled.wrap_units, list(model.blocks), DEFAULT_DTYPE) + for module, policy_dtype in plan: + fully_shard(module, **fsdp_kwargs(policy_dtype)) fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE)) seen: dict[str, torch.dtype] = {} diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 26d826da..ed004dbc 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn -from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision_plan +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, build_wrap_plan, compile_precision_plan class Rope(nn.Module): @@ -103,6 +103,7 @@ def test_empty_module_sel_rejected(): def test_every_node_of_a_nested_chain_wraps_bottom_up(): + model = _model() spec = PrecisionSpec( rules=( Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), @@ -110,11 +111,19 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): Rule(ModuleSel(fqn="blocks.0.attn.norm_q"), gather="default"), ) ) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) - assert _units(compiled) == [ - ("blocks.0.attn.norm_q", torch.bfloat16), - ("blocks.0.attn", torch.float32), - ("blocks.0", torch.float16), + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert dict(_units(compiled)) == { + "blocks.0.attn.norm_q": torch.bfloat16, + "blocks.0.attn": torch.float32, + "blocks.0": torch.float16, + } + # The block that is also a unit keeps its pin, and children wrap before parents. + plan = build_wrap_plan(model, compiled.wrap_units, list(model.blocks), torch.bfloat16) + assert plan == [ + (model.blocks[0].attn.norm_q, torch.bfloat16), + (model.blocks[0].attn, torch.float32), + (model.blocks[0], torch.float16), + (model.blocks[1], torch.bfloat16), ] From 331afed7690a46adc8e79c6964dddf58e30fc448 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 21:59:17 -0700 Subject: [PATCH 22/50] fix(fsdp): wrap block modules at their effective gather dtype --- miles/backends/fsdp_utils/actor.py | 8 +++--- miles/backends/fsdp_utils/precision.py | 28 +++++++++++-------- .../fsdp_utils/_precision_wrap_worker.py | 3 +- .../fsdp_utils/test_precision_plan.py | 16 ++++++++++- 4 files changed, 36 insertions(+), 19 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 4712cfe0..e2270ab7 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -160,7 +160,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), - precision_wrap_units=compiled_precision.wrap_units, + compiled_precision=compiled_precision, ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -630,7 +630,7 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, precision_wrap_units=None): +def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, compiled_precision=None): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard offload_policy = CPUOffloadPolicy() if cpu_offload else None @@ -644,7 +644,7 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules reduce_dtype = resolve_dtype(args.fsdp_reduce_dtype) logger.info( f"FSDP: wrapping {len(modules)} modules of type {layer_cls_to_wrap}, param_dtype={param_dtype}, " - f"reduce_dtype={reduce_dtype}, precision wrap units={len(precision_wrap_units) if precision_wrap_units else 0}" + f"reduce_dtype={reduce_dtype}, precision wrap units={len(compiled_precision.wrap_units)}" ) def _fsdp_kwargs(policy_param_dtype): @@ -659,7 +659,7 @@ def _fsdp_kwargs(policy_param_dtype): "mesh": mesh, } - for module, policy_dtype in build_wrap_plan(model, precision_wrap_units or [], modules, param_dtype): + for module, policy_dtype in build_wrap_plan(model, compiled_precision, modules): fully_shard(module, **_fsdp_kwargs(policy_dtype)) fully_shard(model, **_fsdp_kwargs(param_dtype)) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 504032b0..fbcf80ec 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -115,6 +115,8 @@ class WrapUnit: class CompiledPrecision: master_casts: list[MasterCast] wrap_units: list[WrapUnit] + # Effective gather dtype of every module, i.e. the dtype its innermost wrap unit provides. + gather_dtypes: dict[str, torch.dtype] def apply_master_casts(self) -> None: for cast in self.master_casts: @@ -187,10 +189,10 @@ def compile_precision_plan( # (2)-(4) named_modules is parent-first, so the enclosing unit's dtype is already known here. master_casts: list[MasterCast] = [] wrap_units: list[WrapUnit] = [] - unit_dtypes: dict[str, torch.dtype] = {"": default_dtype} + gather_dtypes: dict[str, torch.dtype] = {"": default_dtype} for mod_fqn, module in model.named_modules(): - enclosing_dtype = unit_dtypes[_parent_fqn(mod_fqn)] - unit_dtypes[mod_fqn] = enclosing_dtype + enclosing_dtype = gather_dtypes[_parent_fqn(mod_fqn)] + gather_dtypes[mod_fqn] = enclosing_dtype master, gather = _fold_covering_rules(spec.rules, matched_fqns, mod_fqn) master_dtype = _resolve_axis(master, default_dtype) @@ -209,23 +211,25 @@ def compile_precision_plan( if mod_fqn == "": raise ValueError("cannot wrap the root module for a gather override") wrap_units.append(WrapUnit(mod_fqn, module, gather_dtype)) - unit_dtypes[mod_fqn] = gather_dtype + gather_dtypes[mod_fqn] = gather_dtype - return CompiledPrecision(master_casts=master_casts, wrap_units=wrap_units) + return CompiledPrecision(master_casts=master_casts, wrap_units=wrap_units, gather_dtypes=gather_dtypes) def build_wrap_plan( model: torch.nn.Module, - wrap_units: list[WrapUnit], + compiled: CompiledPrecision, block_modules: list[torch.nn.Module], - default_dtype: torch.dtype, ) -> list[tuple[torch.nn.Module, torch.dtype]]: - """One wrap order for FSDP2, deepest module first: precision units pin their param_dtype, - remaining block modules take the default, and a module appearing in both keeps the pin.""" - plan: dict[torch.nn.Module, torch.dtype] = {unit.module: unit.param_dtype for unit in wrap_units} + """One wrap order for FSDP2, deepest module first. Block modules are extra wraps that FSDP + needs for sharding granularity, so they must carry their own effective dtype: wrapping one at + the default inside an overridden region would be the innermost wrap and undo the override.""" + plan: dict[torch.nn.Module, torch.dtype] = {unit.module: unit.param_dtype for unit in compiled.wrap_units} + depths, fqns = {}, {} + for mod_fqn, module in model.named_modules(): + depths[module], fqns[module] = mod_fqn.count("."), mod_fqn for module in block_modules: - plan.setdefault(module, default_dtype) - depths = {module: mod_fqn.count(".") for mod_fqn, module in model.named_modules()} + plan.setdefault(module, compiled.gather_dtypes[fqns[module]]) return [(module, plan[module]) for module in sorted(plan, key=lambda module: -depths[module])] diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index 068ae507..802dddca 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -94,8 +94,7 @@ def fsdp_kwargs(param_dtype): policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) return {"mp_policy": policy, "mesh": mesh} - plan = build_wrap_plan(model, compiled.wrap_units, list(model.blocks), DEFAULT_DTYPE) - for module, policy_dtype in plan: + for module, policy_dtype in build_wrap_plan(model, compiled, list(model.blocks)): fully_shard(module, **fsdp_kwargs(policy_dtype)) fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE)) diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index ed004dbc..64adc1fd 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -118,7 +118,7 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): "blocks.0": torch.float16, } # The block that is also a unit keeps its pin, and children wrap before parents. - plan = build_wrap_plan(model, compiled.wrap_units, list(model.blocks), torch.bfloat16) + plan = build_wrap_plan(model, compiled, list(model.blocks)) assert plan == [ (model.blocks[0].attn.norm_q, torch.bfloat16), (model.blocks[0].attn, torch.float32), @@ -127,6 +127,20 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): ] +def test_block_inside_an_override_wraps_at_the_override_dtype(): + model = _model() + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks"), gather="fp32"),)) + compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + assert _units(compiled) == [("blocks", torch.float32)] + # Blocks wrap deeper than the override, so at the default they would undo it. + plan = build_wrap_plan(model, compiled, list(model.blocks)) + assert plan == [ + (model.blocks[0], torch.float32), + (model.blocks[1], torch.float32), + (model.blocks, torch.float32), + ] + + def test_inherited_gather_needs_no_extra_unit(): # attn owns no parameter but its subtree does, so it wraps; norm_q inherits the same dtype. spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Attn"), gather="fp32"),)) From a9c21e1090a51a29d220c2410ba21091d507feec Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 22:18:59 -0700 Subject: [PATCH 23/50] feat(fsdp): keep precision wrap units replicated instead of sharded --- miles/backends/fsdp_utils/actor.py | 21 ++++-- miles/backends/fsdp_utils/parallel.py | 5 +- miles/backends/fsdp_utils/precision.py | 16 ++-- .../fsdp_utils/_precision_wrap_worker.py | 73 ++++++++++++------- .../fsdp_utils/test_precision_plan.py | 69 ++++++++++++++---- 5 files changed, 131 insertions(+), 53 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index e2270ab7..32381df4 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -157,6 +157,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty model = apply_fsdp2( model, mesh=self.parallel_state.get_mesh("fsdp"), + replicate_mesh=self.parallel_state.get_mesh("fsdp_replicate"), cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), @@ -630,7 +631,15 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - return model -def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules=None, compiled_precision=None): +def apply_fsdp2( + model, + mesh=None, + replicate_mesh=None, + cpu_offload=False, + args=None, + no_split_modules=None, + compiled_precision=None, +): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard offload_policy = CPUOffloadPolicy() if cpu_offload else None @@ -647,7 +656,7 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules f"reduce_dtype={reduce_dtype}, precision wrap units={len(compiled_precision.wrap_units)}" ) - def _fsdp_kwargs(policy_param_dtype): + def _fsdp_kwargs(policy_param_dtype, unit_mesh): return { # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent. "mp_policy": MixedPrecisionPolicy( @@ -656,12 +665,12 @@ def _fsdp_kwargs(policy_param_dtype): cast_forward_inputs=False, ), "offload_policy": offload_policy, - "mesh": mesh, + "mesh": unit_mesh, } - for module, policy_dtype in build_wrap_plan(model, compiled_precision, modules): - fully_shard(module, **_fsdp_kwargs(policy_dtype)) + for unit in build_wrap_plan(model, compiled_precision, modules): + fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype, mesh if unit.shard else replicate_mesh)) - fully_shard(model, **_fsdp_kwargs(param_dtype)) + fully_shard(model, **_fsdp_kwargs(param_dtype, mesh)) return model diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 16a113e5..4cc515fa 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -23,7 +23,10 @@ def build_fsdp_meshes( shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "fsdp")) # A degree-1 replicate axis would all-reduce over a single rank every bucket. fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["fsdp"] - meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "dp": world_mesh} + # Degree-1 shard axis: params stay replicated (no all-gather) and grads all-reduce, for the + # small modules the precision spec pins (see precision.WrapUnit.shard). + replicate_mesh = init_device_mesh(device_type, (world_size, 1), mesh_dim_names=("replicate", "shard")) + meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "fsdp_replicate": replicate_mesh, "dp": world_mesh} if sp_size > 1: dp_sp_view = world_mesh._unflatten(0, (world_size // sp_size, sp_size), ("dp", "sp")) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index fbcf80ec..31e35aad 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -32,7 +32,8 @@ modules into one deepest-first order, so FSDP2 always nests child-before-parent — that is how gather="default" carves a module back out of a - non-default ancestor. + non-default ancestor. Precision units wrap on a + replicated mesh (no all-gather, see WrapUnit.shard). Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and FSDP2 requires uniform master dtype among trainable params per unit), so @@ -104,11 +105,13 @@ class MasterCast: @dataclass(frozen=True) class WrapUnit: - """A module to fully_shard on its own with param_dtype=gather.""" + """A module to fully_shard on its own with param_dtype=gather. Precision units are replicated + (shard=False) since a pinned module is small by nature; that will become a size-driven choice.""" fqn: str module: torch.nn.Module param_dtype: torch.dtype + shard: bool = False @dataclass @@ -220,17 +223,18 @@ def build_wrap_plan( model: torch.nn.Module, compiled: CompiledPrecision, block_modules: list[torch.nn.Module], -) -> list[tuple[torch.nn.Module, torch.dtype]]: +) -> list[WrapUnit]: """One wrap order for FSDP2, deepest module first. Block modules are extra wraps that FSDP needs for sharding granularity, so they must carry their own effective dtype: wrapping one at the default inside an overridden region would be the innermost wrap and undo the override.""" - plan: dict[torch.nn.Module, torch.dtype] = {unit.module: unit.param_dtype for unit in compiled.wrap_units} + plan: dict[torch.nn.Module, WrapUnit] = {unit.module: unit for unit in compiled.wrap_units} depths, fqns = {}, {} for mod_fqn, module in model.named_modules(): depths[module], fqns[module] = mod_fqn.count("."), mod_fqn for module in block_modules: - plan.setdefault(module, compiled.gather_dtypes[fqns[module]]) - return [(module, plan[module]) for module in sorted(plan, key=lambda module: -depths[module])] + fqn = fqns[module] + plan.setdefault(module, WrapUnit(fqn, module, compiled.gather_dtypes[fqn], shard=True)) + return [plan[module] for module in sorted(plan, key=lambda module: -depths[module])] def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index 802dddca..a36a951a 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -1,7 +1,20 @@ """Gloo worker asserting the compiled precision plan really wraps under FSDP2 (2 ranks). +Model and spec (default gather dtype is bf16, master is fp32): + + Net + ├── stem Leaf -> bf16, sharded (no rule) + └── blocks + ├── 0: Block -> block unit, bf16, sharded + │ ├── norm Norm -> fp32 via cls rule, replicated + │ └── attn Attn -> fp16 via fqn rule, replicated + │ ├── norm_q Norm -> bf16, carved back out of attn + │ └── proj Leaf -> fp16, inherits attn + └── 1: Block -> same, except norm_q keeps fp32 + Modules cast explicitly in forward because CPU kernels reject mixed dtypes; what -matters here is the wrap nesting and the param dtype each module sees at forward. +matters here is the wrap nesting, the param dtype each module sees at forward, +and that precision units are replicated instead of sharded. """ import torch @@ -9,13 +22,14 @@ import torch.nn as nn from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard +from torch.distributed.tensor import DTensor, Replicate from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, build_wrap_plan, compile_precision_plan DEFAULT_DTYPE = torch.bfloat16 -class Norm(nn.Module): +class Leaf(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.ones(8)) @@ -24,20 +38,15 @@ def forward(self, x): return x * self.weight.to(x.dtype) -class Proj(nn.Module): - def __init__(self): - super().__init__() - self.weight = nn.Parameter(torch.eye(8)) - - def forward(self, x): - return x @ self.weight.to(x.dtype) +class Norm(Leaf): + pass class Attn(nn.Module): def __init__(self): super().__init__() self.norm_q = Norm() - self.proj = Proj() + self.proj = Leaf() def forward(self, x): return self.proj(self.norm_q(x)) @@ -53,12 +62,14 @@ def forward(self, x): return self.attn(self.norm(x)) -class Tiny(nn.Module): +class Net(nn.Module): def __init__(self): super().__init__() + self.stem = Leaf() self.blocks = nn.ModuleList([Block(), Block()]) def forward(self, x): + x = self.stem(x) for block in self.blocks: x = block(x) return x @@ -66,37 +77,47 @@ def forward(self, x): SPEC = PrecisionSpec( rules=( - Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), - Rule(ModuleSel(fqn="blocks.0.attn"), gather="fp32"), + Rule(ModuleSel(cls="Norm"), gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.attn"), gather="fp16"), Rule(ModuleSel(fqn="blocks.0.attn.norm_q"), gather="default"), - Rule(ModuleSel(cls="Norm", fqn="blocks.1*"), gather="fp32"), ) ) EXPECTED_GATHER = { - "blocks.0.norm": torch.float16, # inherits the fp16 block unit - "blocks.0.attn.norm_q": DEFAULT_DTYPE, # carved back out of two non-default ancestors - "blocks.0.attn.proj": torch.float32, # inherits the fp32 attn unit - "blocks.1.norm": torch.float32, # cls + fqn rule - "blocks.1.attn.norm_q": torch.float32, - "blocks.1.attn.proj": DEFAULT_DTYPE, # untouched by any rule + "stem": DEFAULT_DTYPE, # no rule + "blocks.0.norm": torch.float32, # cls rule + "blocks.0.attn.norm_q": DEFAULT_DTYPE, # carved back out of the fp16 attn unit + "blocks.0.attn.proj": torch.float16, # inherits the fp16 attn unit + "blocks.1.norm": torch.float32, # cls rule + "blocks.1.attn.norm_q": torch.float32, # cls rule, no override above it + "blocks.1.attn.proj": DEFAULT_DTYPE, # no rule } def main() -> None: dist.init_process_group("gloo") - mesh = init_device_mesh("cpu", (dist.get_world_size(),)) - model = Tiny().to(torch.float32) # fp32 master + world_size = dist.get_world_size() + shard_mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("fsdp",)) + replicate_mesh = init_device_mesh("cpu", (world_size, 1), mesh_dim_names=("replicate", "shard")) + model = Net().to(torch.float32) # fp32 master compiled = compile_precision_plan(model, SPEC, default_dtype=DEFAULT_DTYPE) compiled.apply_master_casts() - def fsdp_kwargs(param_dtype): + def fsdp_kwargs(param_dtype, mesh): policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) return {"mp_policy": policy, "mesh": mesh} - for module, policy_dtype in build_wrap_plan(model, compiled, list(model.blocks)): - fully_shard(module, **fsdp_kwargs(policy_dtype)) - fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE)) + for unit in build_wrap_plan(model, compiled, list(model.blocks)): + fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype, shard_mesh if unit.shard else replicate_mesh)) + fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE, shard_mesh)) + + # Precision units are replicated, so their local shard is the whole tensor: no all-gather. + for unit in compiled.wrap_units: + param = next(unit.module.parameters()) + if not isinstance(param, DTensor) or param.placements[0] != Replicate(): + raise AssertionError(f"{unit.fqn} is not replicated: {getattr(param, 'placements', None)}") + if param.to_local().shape != param.shape: + raise AssertionError(f"{unit.fqn} local shard {param.to_local().shape} != full {param.shape}") seen: dict[str, torch.dtype] = {} diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 64adc1fd..12687cca 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -1,3 +1,18 @@ +"""Compiling PrecisionSpec rules into master casts and FSDP2 wrap units. + +Test model (every test uses two identical blocks; `w` marks own float params): + + Tiny + └── blocks + ├── 0: Block + │ ├── linear Linear w + │ ├── norm LayerNorm w + │ ├── attn Attn (no own param) + │ │ └── norm_q LayerNorm w + │ └── rope Rope (buffer only) + └── 1: Block (same) +""" + from tests.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[]) @@ -44,16 +59,23 @@ def _units(compiled): return [(unit.fqn, unit.param_dtype) for unit in compiled.wrap_units] +def _plan(model, compiled): + return [(unit.fqn, unit.param_dtype, unit.shard) for unit in build_wrap_plan(model, compiled, list(model.blocks))] + + NORM_FQNS = {f"blocks.{i}{suffix}" for i in range(2) for suffix in (".norm", ".attn.norm_q")} def test_empty_spec_compiles_to_nothing(): + """No rules -> the model keeps the run's default dtype everywhere.""" compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert compiled.wrap_units == [] def test_fqn_glob_selects_norms_across_depths(): + """`*norm*` crosses dots, so one rule catches `blocks.N.norm` (depth 2) and + `blocks.N.attn.norm_q` (depth 3) while leaving their Linear siblings alone.""" model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -64,6 +86,7 @@ def test_fqn_glob_selects_norms_across_depths(): def test_cls_glob_selects_by_class(): + """Selecting by class name reaches the same norms without naming any path.""" spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="*LayerNorm"), gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] @@ -71,6 +94,8 @@ def test_cls_glob_selects_by_class(): def test_master_rule_covers_matched_subtree(): + """A rule on `blocks.1` casts every float tensor under it (params and buffers), + and nothing under its sibling `blocks.0`.""" model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -81,6 +106,11 @@ def test_master_rule_covers_matched_subtree(): def test_later_rule_overrides_earlier_selection(): + """Rules overlap on block 0's norms; the later rule wins there, block 1 keeps the first rule. + + blocks.0.norm, blocks.0.attn.norm_q -> fp16 (rule 2) + blocks.1.norm, blocks.1.attn.norm_q -> fp32 (rule 1) + """ spec = PrecisionSpec( rules=( Rule(ModuleSel(cls="LayerNorm"), gather="fp32"), @@ -97,12 +127,21 @@ def test_later_rule_overrides_earlier_selection(): def test_empty_module_sel_rejected(): + """A selector with neither fqn nor cls would match everything by accident.""" spec = PrecisionSpec(rules=(Rule(ModuleSel(), master="fp32"),)) with pytest.raises(ValueError, match="empty ModuleSel"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) def test_every_node_of_a_nested_chain_wraps_bottom_up(): + """Three nested rules, each differing from its parent, need one unit per node: + + blocks.0 fp16 <- wraps last (shallowest) + └── attn fp32 + └── norm_q bf16 <- wraps first, so the outer units exclude it + + `blocks.1` is only wrapped as a block unit, at the default dtype. + """ model = _model() spec = PrecisionSpec( rules=( @@ -117,38 +156,39 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): "blocks.0.attn": torch.float32, "blocks.0": torch.float16, } - # The block that is also a unit keeps its pin, and children wrap before parents. - plan = build_wrap_plan(model, compiled, list(model.blocks)) - assert plan == [ - (model.blocks[0].attn.norm_q, torch.bfloat16), - (model.blocks[0].attn, torch.float32), - (model.blocks[0], torch.float16), - (model.blocks[1], torch.bfloat16), + # Precision units are replicated (shard=False); the plain block unit still shards. + assert _plan(model, compiled) == [ + ("blocks.0.attn.norm_q", torch.bfloat16, False), + ("blocks.0.attn", torch.float32, False), + ("blocks.0", torch.float16, False), + ("blocks.1", torch.bfloat16, True), ] def test_block_inside_an_override_wraps_at_the_override_dtype(): + """`blocks` (the ModuleList) is an ancestor of both block units, so the blocks wrap deeper + than the override; at the default dtype they would be the innermost wrap and undo it.""" model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks"), gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert _units(compiled) == [("blocks", torch.float32)] - # Blocks wrap deeper than the override, so at the default they would undo it. - plan = build_wrap_plan(model, compiled, list(model.blocks)) - assert plan == [ - (model.blocks[0], torch.float32), - (model.blocks[1], torch.float32), - (model.blocks, torch.float32), + assert _plan(model, compiled) == [ + ("blocks.0", torch.float32, True), + ("blocks.1", torch.float32, True), + ("blocks", torch.float32, False), ] def test_inherited_gather_needs_no_extra_unit(): - # attn owns no parameter but its subtree does, so it wraps; norm_q inherits the same dtype. + """`attn` owns no parameter but its subtree does, so it wraps; `norm_q` inherits the same + dtype from it and needs no unit of its own (the units are a minimal cover of the tree).""" spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Attn"), gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert _units(compiled) == [("blocks.0.attn", torch.float32), ("blocks.1.attn", torch.float32)] def test_buffer_only_module_casts_master_without_wrapping(): + """Buffers are never gathered, so a paramless module takes the master cast and no unit.""" model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -157,6 +197,7 @@ def test_buffer_only_module_casts_master_without_wrapping(): def test_unmatched_rule_rejected(): + """A rule matching nothing is a typo'd pattern or class name, not a no-op.""" spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) with pytest.raises(ValueError, match="matched no module"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) From b854a385f9c8d400c652202e04f4f89ce02a9f28 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 22:20:38 -0700 Subject: [PATCH 24/50] docs(test): draw the module tree and per-node dtypes in the precision tests --- .../fsdp_utils/_precision_wrap_worker.py | 29 ++-- .../fsdp_utils/test_precision_plan.py | 144 +++++++++++++----- 2 files changed, 129 insertions(+), 44 deletions(-) diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index a36a951a..b5716a10 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -1,16 +1,27 @@ """Gloo worker asserting the compiled precision plan really wraps under FSDP2 (2 ranks). -Model and spec (default gather dtype is bf16, master is fp32): +Master is fp32 everywhere, default gather dtype is bf16, and the spec is - Net - ├── stem Leaf -> bf16, sharded (no rule) + Rule(cls="Norm", gather=fp32) + Rule(fqn="blocks.0.attn", gather=fp16) + Rule(fqn="blocks.0.attn.norm_q", gather=default) + +so the tree, the dtype each module's params carry in the forward, and how each +wrap unit is placed on the mesh come out as: + + Net gather placement + ├── stem Leaf bf16 sharded (root unit, no rule) └── blocks - ├── 0: Block -> block unit, bf16, sharded - │ ├── norm Norm -> fp32 via cls rule, replicated - │ └── attn Attn -> fp16 via fqn rule, replicated - │ ├── norm_q Norm -> bf16, carved back out of attn - │ └── proj Leaf -> fp16, inherits attn - └── 1: Block -> same, except norm_q keeps fp32 + ├── 0 Block [U] bf16 sharded (block unit) + │ ├── norm Norm [U] fp32 replicated + │ └── attn Attn [U] fp16 replicated + │ ├── norm_q Norm [U] bf16 replicated (carved out of attn) + │ └── proj Leaf fp16 (inside the attn unit) + └── 1 Block [U] bf16 sharded (block unit) + ├── norm Norm [U] fp32 replicated + └── attn Attn + ├── norm_q Norm [U] fp32 replicated (no override above it) + └── proj Leaf bf16 (inside the block unit) Modules cast explicitly in forward because CPU kernels reject mixed dtypes; what matters here is the wrap nesting, the param dtype each module sees at forward, diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 12687cca..a91aa758 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -1,16 +1,18 @@ """Compiling PrecisionSpec rules into master casts and FSDP2 wrap units. -Test model (every test uses two identical blocks; `w` marks own float params): - - Tiny - └── blocks - ├── 0: Block - │ ├── linear Linear w - │ ├── norm LayerNorm w - │ ├── attn Attn (no own param) - │ │ └── norm_q LayerNorm w - │ └── rope Rope (buffer only) - └── 1: Block (same) +Every test uses this model, loaded at bf16 with default_dtype=bf16, and every +docstring draws the resulting dtype per node (`[U]` = its own wrap unit). +`blocks.1` mirrors `blocks.0`, so most diagrams only draw block 0. + + Tiny classes and own float tensors + └── blocks ModuleList - + ├── 0 Block - + │ ├── linear Linear weight, bias + │ ├── norm LayerNorm weight, bias + │ ├── attn Attn - + │ │ └── norm_q LayerNorm weight, bias + │ └── rope Rope freqs (buffer only) + └── 1 Block (same) """ from tests.ci.ci_register import register_cpu_ci @@ -67,15 +69,32 @@ def _plan(model, compiled): def test_empty_spec_compiles_to_nothing(): - """No rules -> the model keeps the run's default dtype everywhere.""" + """No rules, so every node keeps the default and nothing is emitted. + + blocks.0 gather bf16 + ├── linear bf16 + ├── norm bf16 + ├── attn bf16 + │ └── norm_q bf16 + └── rope bf16 + """ compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert compiled.wrap_units == [] def test_fqn_glob_selects_norms_across_depths(): - """`*norm*` crosses dots, so one rule catches `blocks.N.norm` (depth 2) and - `blocks.N.attn.norm_q` (depth 3) while leaving their Linear siblings alone.""" + """`*norm*` crosses dots, so one rule catches both norm depths and skips the Linear siblings. + + Rule(fqn="*norm*", master=fp32, gather=fp32) + + blocks.0 master bf16 gather bf16 + ├── linear bf16 bf16 + ├── norm [U] fp32 fp32 + ├── attn bf16 bf16 + │ └── norm_q [U] fp32 fp32 + └── rope bf16 bf16 + """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -86,7 +105,17 @@ def test_fqn_glob_selects_norms_across_depths(): def test_cls_glob_selects_by_class(): - """Selecting by class name reaches the same norms without naming any path.""" + """Selecting by class name reaches the same two norms without naming any path. + + Rule(cls="*LayerNorm", gather=fp32) + + blocks.0 gather bf16 + ├── linear bf16 + ├── norm [U] fp32 + ├── attn bf16 + │ └── norm_q [U] fp32 + └── rope bf16 + """ spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="*LayerNorm"), gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] @@ -94,8 +123,17 @@ def test_cls_glob_selects_by_class(): def test_master_rule_covers_matched_subtree(): - """A rule on `blocks.1` casts every float tensor under it (params and buffers), - and nothing under its sibling `blocks.0`.""" + """A rule covers its module's whole subtree, params and buffers alike, and stops at the sibling. + + Rule(fqn="blocks.1", master=fp32) + + blocks.0 master bf16 blocks.1 master fp32 + ├── linear bf16 ├── linear fp32 + ├── norm bf16 ├── norm fp32 + ├── attn bf16 ├── attn fp32 + │ └── norm_q bf16 │ └── norm_q fp32 + └── rope.freqs bf16 └── rope.freqs fp32 + """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -106,10 +144,17 @@ def test_master_rule_covers_matched_subtree(): def test_later_rule_overrides_earlier_selection(): - """Rules overlap on block 0's norms; the later rule wins there, block 1 keeps the first rule. + """Both rules cover block 0's norms; the later one wins there while block 1 keeps the first. + + Rule(cls="LayerNorm", gather=fp32) # rule 1 + Rule(fqn="blocks.0.*norm*", gather=fp16) # rule 2, wins where they overlap - blocks.0.norm, blocks.0.attn.norm_q -> fp16 (rule 2) - blocks.1.norm, blocks.1.attn.norm_q -> fp32 (rule 1) + blocks.0 gather bf16 blocks.1 gather bf16 + ├── linear bf16 ├── linear bf16 + ├── norm [U] fp16 (rule 2) ├── norm [U] fp32 (rule 1) + ├── attn bf16 ├── attn bf16 + │ └── norm_q [U] fp16 (rule 2) │ └── norm_q [U] fp32 (rule 1) + └── rope bf16 └── rope bf16 """ spec = PrecisionSpec( rules=( @@ -127,20 +172,27 @@ def test_later_rule_overrides_earlier_selection(): def test_empty_module_sel_rejected(): - """A selector with neither fqn nor cls would match everything by accident.""" + """A selector with neither fqn nor cls would silently match every module.""" spec = PrecisionSpec(rules=(Rule(ModuleSel(), master="fp32"),)) with pytest.raises(ValueError, match="empty ModuleSel"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) def test_every_node_of_a_nested_chain_wraps_bottom_up(): - """Three nested rules, each differing from its parent, need one unit per node: - - blocks.0 fp16 <- wraps last (shallowest) - └── attn fp32 - └── norm_q bf16 <- wraps first, so the outer units exclude it - - `blocks.1` is only wrapped as a block unit, at the default dtype. + """Three nested rules that each differ from their parent need one unit per node, and the plan + hands them back deepest first so each outer wrap excludes the inner ones. + + Rule(fqn="blocks.0", gather=fp16) + Rule(fqn="blocks.0.attn", gather=fp32) + Rule(fqn="blocks.0.attn.norm_q", gather=default) # carved back out + + blocks.0 [U] fp16 gather order 3 + ├── linear fp16 (inherits blocks.0) + ├── norm fp16 (inherits blocks.0) + ├── attn [U] fp32 gather order 2 + │ └── norm_q [U] bf16 gather order 1, wraps first + └── rope fp16 buffer, never gathered + blocks.1 [U] bf16 block unit only, still sharded """ model = _model() spec = PrecisionSpec( @@ -156,7 +208,7 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): "blocks.0.attn": torch.float32, "blocks.0": torch.float16, } - # Precision units are replicated (shard=False); the plain block unit still shards. + # (fqn, param_dtype, shard): precision units are replicated, the plain block unit shards. assert _plan(model, compiled) == [ ("blocks.0.attn.norm_q", torch.bfloat16, False), ("blocks.0.attn", torch.float32, False), @@ -166,8 +218,15 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): def test_block_inside_an_override_wraps_at_the_override_dtype(): - """`blocks` (the ModuleList) is an ancestor of both block units, so the blocks wrap deeper - than the override; at the default dtype they would be the innermost wrap and undo it.""" + """The rule sits above the block units, so the blocks wrap deeper than the override; at the + default dtype they would be the innermost wrap and silently undo it. + + Rule(fqn="blocks", gather=fp32) + + blocks [U] fp32 gather order 3 (the override) + ├── 0 fp32 gather order 1, block unit forced to fp32 + └── 1 fp32 gather order 2, block unit forced to fp32 + """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks"), gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -180,15 +239,30 @@ def test_block_inside_an_override_wraps_at_the_override_dtype(): def test_inherited_gather_needs_no_extra_unit(): - """`attn` owns no parameter but its subtree does, so it wraps; `norm_q` inherits the same - dtype from it and needs no unit of its own (the units are a minimal cover of the tree).""" + """Units are a minimal cover: a node whose dtype already comes from an enclosing unit is skipped. + + Rule(cls="Attn", gather=fp32) + + blocks.0 gather bf16 + ├── linear bf16 + ├── norm bf16 + ├── attn [U] fp32 (owns no param, but its subtree does) + │ └── norm_q fp32 inherits attn, so no unit of its own + └── rope bf16 + """ spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Attn"), gather="fp32"),)) compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) assert _units(compiled) == [("blocks.0.attn", torch.float32), ("blocks.1.attn", torch.float32)] def test_buffer_only_module_casts_master_without_wrapping(): - """Buffers are never gathered, so a paramless module takes the master cast and no unit.""" + """Buffers are never gathered, so a paramless module takes the master cast and no unit. + + Rule(cls="Rope", master=fp32, gather=fp32) + + blocks.0 master bf16 gather bf16 + └── rope.freqs fp32 n/a (buffer), no unit + """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32", gather="fp32"),)) compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) @@ -197,7 +271,7 @@ def test_buffer_only_module_casts_master_without_wrapping(): def test_unmatched_rule_rejected(): - """A rule matching nothing is a typo'd pattern or class name, not a no-op.""" + """A rule matching nothing is a typo'd pattern or class name, not a silent no-op.""" spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) with pytest.raises(ValueError, match="matched no module"): compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) From 3eacf33044f5a0986e49b46063957f7e61751cc0 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 22:48:28 -0700 Subject: [PATCH 25/50] refactor(fsdp): derive the no-shard mesh from world_mesh like the other views --- miles/backends/fsdp_utils/actor.py | 6 +++--- miles/backends/fsdp_utils/parallel.py | 8 ++++---- tests/fast/backends/fsdp_utils/_precision_wrap_worker.py | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 32381df4..188d39df 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -157,7 +157,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty model = apply_fsdp2( model, mesh=self.parallel_state.get_mesh("fsdp"), - replicate_mesh=self.parallel_state.get_mesh("fsdp_replicate"), + noshard_mesh=self.parallel_state.get_mesh("fsdp_noshard"), cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), @@ -634,7 +634,7 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - def apply_fsdp2( model, mesh=None, - replicate_mesh=None, + noshard_mesh=None, cpu_offload=False, args=None, no_split_modules=None, @@ -669,7 +669,7 @@ def _fsdp_kwargs(policy_param_dtype, unit_mesh): } for unit in build_wrap_plan(model, compiled_precision, modules): - fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype, mesh if unit.shard else replicate_mesh)) + fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype, mesh if unit.shard else noshard_mesh)) fully_shard(model, **_fsdp_kwargs(param_dtype, mesh)) diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 4cc515fa..d7a31bff 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -23,10 +23,10 @@ def build_fsdp_meshes( shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "fsdp")) # A degree-1 replicate axis would all-reduce over a single rank every bucket. fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["fsdp"] - # Degree-1 shard axis: params stay replicated (no all-gather) and grads all-reduce, for the - # small modules the precision spec pins (see precision.WrapUnit.shard). - replicate_mesh = init_device_mesh(device_type, (world_size, 1), mesh_dim_names=("replicate", "shard")) - meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "fsdp_replicate": replicate_mesh, "dp": world_mesh} + # Same ranks with a degree-1 shard axis: params stay replicated (no all-gather) and grads + # all-reduce, for the small modules the precision spec pins (see precision.WrapUnit.shard). + noshard_mesh = world_mesh._unflatten(0, (world_size, 1), ("dp_replicate", "dp_shard")) + meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "fsdp_noshard": noshard_mesh, "dp": world_mesh} if sp_size > 1: dp_sp_view = world_mesh._unflatten(0, (world_size // sp_size, sp_size), ("dp", "sp")) diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index b5716a10..ce8ff10c 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -108,7 +108,7 @@ def main() -> None: dist.init_process_group("gloo") world_size = dist.get_world_size() shard_mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("fsdp",)) - replicate_mesh = init_device_mesh("cpu", (world_size, 1), mesh_dim_names=("replicate", "shard")) + noshard_mesh = init_device_mesh("cpu", (world_size, 1), mesh_dim_names=("dp_replicate", "dp_shard")) model = Net().to(torch.float32) # fp32 master compiled = compile_precision_plan(model, SPEC, default_dtype=DEFAULT_DTYPE) @@ -119,7 +119,7 @@ def fsdp_kwargs(param_dtype, mesh): return {"mp_policy": policy, "mesh": mesh} for unit in build_wrap_plan(model, compiled, list(model.blocks)): - fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype, shard_mesh if unit.shard else replicate_mesh)) + fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype, shard_mesh if unit.shard else noshard_mesh)) fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE, shard_mesh)) # Precision units are replicated, so their local shard is the whole tensor: no all-gather. From 8f85203c45fe883d5aeafc84b725588f4b0b123e Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 22:57:03 -0700 Subject: [PATCH 26/50] refactor(fsdp): name the shard axis dp_shard and assert the no-shard view --- miles/backends/fsdp_utils/parallel.py | 8 ++++---- .../fsdp_utils/_hybrid_shard_mesh_worker.py | 13 ++++++++++--- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index d7a31bff..65eb196b 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -20,11 +20,11 @@ def build_fsdp_meshes( """Build the FSDP hybrid-shard and DP/SP views.""" world_mesh = init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",)) - shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "fsdp")) + shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "dp_shard")) # A degree-1 replicate axis would all-reduce over a single rank every bucket. - fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["fsdp"] - # Same ranks with a degree-1 shard axis: params stay replicated (no all-gather) and grads - # all-reduce, for the small modules the precision spec pins (see precision.WrapUnit.shard). + fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["dp_shard"] + # Same view with dp_shard degree 1: the shard ranks replicate instead, so FSDP keeps these + # params whole (no all-gather) and reduces their grads with one all-reduce over every rank. noshard_mesh = world_mesh._unflatten(0, (world_size, 1), ("dp_replicate", "dp_shard")) meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "fsdp_noshard": noshard_mesh, "dp": world_mesh} diff --git a/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py b/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py index 667af37a..92c94aa7 100644 --- a/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py +++ b/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py @@ -61,16 +61,23 @@ def check(rank, world_size, dp_replicate, dp_shard, ring_degree, ulysses_degree) fsdp = meshes["fsdp"] if dp_replicate > 1: assert fsdp.ndim == 2, fsdp - assert fsdp.mesh_dim_names == ("dp_replicate", "fsdp") + assert fsdp.mesh_dim_names == ("dp_replicate", "dp_shard") assert fsdp["dp_replicate"].size() == dp_replicate - assert fsdp["fsdp"].size() == dp_shard * sp_size + assert fsdp["dp_shard"].size() == dp_shard * sp_size # dp_replicate outermost makes a shard group a contiguous rank run, i.e. one node. - shard_ranks = mesh_group_ranks(fsdp["fsdp"]) + shard_ranks = mesh_group_ranks(fsdp["dp_shard"]) assert shard_ranks == list(range(shard_ranks[0], shard_ranks[0] + len(shard_ranks))) else: assert fsdp.ndim == 1, fsdp assert fsdp.size() == dp_shard * sp_size + # Precision units wrap here: dp_shard degree 1 means nothing is gathered, and every rank + # takes part in the one grad all-reduce. + noshard = meshes["fsdp_noshard"] + assert noshard.mesh_dim_names == ("dp_replicate", "dp_shard") + assert noshard.shape == (world_size, 1) + assert mesh_group_ranks(noshard["dp_replicate"]) == list(range(world_size)) + def main(): dist.init_process_group("gloo") From 0ba9bd1130c8c6e20613730d10901e365a7f40e0 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 22:59:19 -0700 Subject: [PATCH 27/50] refactor(fsdp): build the no-shard view from the same shard_view expression --- miles/backends/fsdp_utils/parallel.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 65eb196b..61cfa276 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -20,12 +20,16 @@ def build_fsdp_meshes( """Build the FSDP hybrid-shard and DP/SP views.""" world_mesh = init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",)) - shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "dp_shard")) + def shard_view(replicate: int) -> DeviceMesh: + """FSDP shards over the last axis, so replicate=world_size leaves nothing to shard.""" + return world_mesh._unflatten(0, (replicate, world_size // replicate), ("dp_replicate", "dp_shard")) + # A degree-1 replicate axis would all-reduce over a single rank every bucket. - fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["dp_shard"] - # Same view with dp_shard degree 1: the shard ranks replicate instead, so FSDP keeps these - # params whole (no all-gather) and reduces their grads with one all-reduce over every rank. - noshard_mesh = world_mesh._unflatten(0, (world_size, 1), ("dp_replicate", "dp_shard")) + hybrid_view = shard_view(dp_replicate) + fsdp_mesh = hybrid_view if dp_replicate > 1 else hybrid_view["dp_shard"] + # Precision units that must not be gathered wrap on the fully replicated view instead; FSDP + # keeps their params whole and reduces the grads with one all-reduce over every rank. + noshard_mesh = shard_view(world_size) meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "fsdp_noshard": noshard_mesh, "dp": world_mesh} if sp_size > 1: From 39ef3131fe6d77efd3e53528184042949bc72923 Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 23:07:47 -0700 Subject: [PATCH 28/50] revert(fsdp): drop the no-shard mesh and keep every wrap unit sharded --- miles/backends/fsdp_utils/actor.py | 10 ++-- miles/backends/fsdp_utils/parallel.py | 13 ++---- miles/backends/fsdp_utils/precision.py | 9 ++-- .../fsdp_utils/_hybrid_shard_mesh_worker.py | 7 --- .../fsdp_utils/_precision_wrap_worker.py | 46 +++++++------------ .../fsdp_utils/test_precision_plan.py | 19 ++++---- 6 files changed, 36 insertions(+), 68 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 188d39df..77cc677b 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -157,7 +157,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty model = apply_fsdp2( model, mesh=self.parallel_state.get_mesh("fsdp"), - noshard_mesh=self.parallel_state.get_mesh("fsdp_noshard"), cpu_offload=self.args.fsdp_cpu_offload, args=self.args, no_split_modules=self.model_backend.fsdp_no_split_modules(model), @@ -634,7 +633,6 @@ def apply_lora(model: torch.nn.Module, args: Namespace, train_pipeline_config) - def apply_fsdp2( model, mesh=None, - noshard_mesh=None, cpu_offload=False, args=None, no_split_modules=None, @@ -656,7 +654,7 @@ def apply_fsdp2( f"reduce_dtype={reduce_dtype}, precision wrap units={len(compiled_precision.wrap_units)}" ) - def _fsdp_kwargs(policy_param_dtype, unit_mesh): + def _fsdp_kwargs(policy_param_dtype): return { # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent. "mp_policy": MixedPrecisionPolicy( @@ -665,12 +663,12 @@ def _fsdp_kwargs(policy_param_dtype, unit_mesh): cast_forward_inputs=False, ), "offload_policy": offload_policy, - "mesh": unit_mesh, + "mesh": mesh, } for unit in build_wrap_plan(model, compiled_precision, modules): - fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype, mesh if unit.shard else noshard_mesh)) + fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype)) - fully_shard(model, **_fsdp_kwargs(param_dtype, mesh)) + fully_shard(model, **_fsdp_kwargs(param_dtype)) return model diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 61cfa276..f5fda80c 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -20,17 +20,10 @@ def build_fsdp_meshes( """Build the FSDP hybrid-shard and DP/SP views.""" world_mesh = init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",)) - def shard_view(replicate: int) -> DeviceMesh: - """FSDP shards over the last axis, so replicate=world_size leaves nothing to shard.""" - return world_mesh._unflatten(0, (replicate, world_size // replicate), ("dp_replicate", "dp_shard")) - + shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "dp_shard")) # A degree-1 replicate axis would all-reduce over a single rank every bucket. - hybrid_view = shard_view(dp_replicate) - fsdp_mesh = hybrid_view if dp_replicate > 1 else hybrid_view["dp_shard"] - # Precision units that must not be gathered wrap on the fully replicated view instead; FSDP - # keeps their params whole and reduces the grads with one all-reduce over every rank. - noshard_mesh = shard_view(world_size) - meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "fsdp_noshard": noshard_mesh, "dp": world_mesh} + fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["dp_shard"] + meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "dp": world_mesh} if sp_size > 1: dp_sp_view = world_mesh._unflatten(0, (world_size // sp_size, sp_size), ("dp", "sp")) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 31e35aad..be6295d9 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -32,8 +32,7 @@ modules into one deepest-first order, so FSDP2 always nests child-before-parent — that is how gather="default" carves a module back out of a - non-default ancestor. Precision units wrap on a - replicated mesh (no all-gather, see WrapUnit.shard). + non-default ancestor. Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and FSDP2 requires uniform master dtype among trainable params per unit), so @@ -105,13 +104,11 @@ class MasterCast: @dataclass(frozen=True) class WrapUnit: - """A module to fully_shard on its own with param_dtype=gather. Precision units are replicated - (shard=False) since a pinned module is small by nature; that will become a size-driven choice.""" + """A module to fully_shard on its own with param_dtype=gather.""" fqn: str module: torch.nn.Module param_dtype: torch.dtype - shard: bool = False @dataclass @@ -233,7 +230,7 @@ def build_wrap_plan( depths[module], fqns[module] = mod_fqn.count("."), mod_fqn for module in block_modules: fqn = fqns[module] - plan.setdefault(module, WrapUnit(fqn, module, compiled.gather_dtypes[fqn], shard=True)) + plan.setdefault(module, WrapUnit(fqn, module, compiled.gather_dtypes[fqn])) return [plan[module] for module in sorted(plan, key=lambda module: -depths[module])] diff --git a/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py b/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py index 92c94aa7..da1dc15a 100644 --- a/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py +++ b/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py @@ -71,13 +71,6 @@ def check(rank, world_size, dp_replicate, dp_shard, ring_degree, ulysses_degree) assert fsdp.ndim == 1, fsdp assert fsdp.size() == dp_shard * sp_size - # Precision units wrap here: dp_shard degree 1 means nothing is gathered, and every rank - # takes part in the one grad all-reduce. - noshard = meshes["fsdp_noshard"] - assert noshard.mesh_dim_names == ("dp_replicate", "dp_shard") - assert noshard.shape == (world_size, 1) - assert mesh_group_ranks(noshard["dp_replicate"]) == list(range(world_size)) - def main(): dist.init_process_group("gloo") diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index ce8ff10c..d8d34a19 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -6,26 +6,24 @@ Rule(fqn="blocks.0.attn", gather=fp16) Rule(fqn="blocks.0.attn.norm_q", gather=default) -so the tree, the dtype each module's params carry in the forward, and how each -wrap unit is placed on the mesh come out as: +so the tree and the dtype each module's params carry in the forward come out as: - Net gather placement - ├── stem Leaf bf16 sharded (root unit, no rule) + Net gather + ├── stem Leaf bf16 (no rule, wrapped by the root unit) └── blocks - ├── 0 Block [U] bf16 sharded (block unit) - │ ├── norm Norm [U] fp32 replicated - │ └── attn Attn [U] fp16 replicated - │ ├── norm_q Norm [U] bf16 replicated (carved out of attn) - │ └── proj Leaf fp16 (inside the attn unit) - └── 1 Block [U] bf16 sharded (block unit) - ├── norm Norm [U] fp32 replicated + ├── 0 Block [U] bf16 (block unit) + │ ├── norm Norm [U] fp32 + │ └── attn Attn [U] fp16 + │ ├── norm_q Norm [U] bf16 (carved back out of attn) + │ └── proj Leaf fp16 (inside the attn unit) + └── 1 Block [U] bf16 (block unit) + ├── norm Norm [U] fp32 └── attn Attn - ├── norm_q Norm [U] fp32 replicated (no override above it) - └── proj Leaf bf16 (inside the block unit) + ├── norm_q Norm [U] fp32 (no override above it) + └── proj Leaf bf16 (inside the block unit) Modules cast explicitly in forward because CPU kernels reject mixed dtypes; what -matters here is the wrap nesting, the param dtype each module sees at forward, -and that precision units are replicated instead of sharded. +matters here is the wrap nesting and the param dtype each module sees at forward. """ import torch @@ -33,7 +31,6 @@ import torch.nn as nn from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard -from torch.distributed.tensor import DTensor, Replicate from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, build_wrap_plan, compile_precision_plan @@ -107,28 +104,19 @@ def forward(self, x): def main() -> None: dist.init_process_group("gloo") world_size = dist.get_world_size() - shard_mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("fsdp",)) - noshard_mesh = init_device_mesh("cpu", (world_size, 1), mesh_dim_names=("dp_replicate", "dp_shard")) + mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("dp_shard",)) model = Net().to(torch.float32) # fp32 master compiled = compile_precision_plan(model, SPEC, default_dtype=DEFAULT_DTYPE) compiled.apply_master_casts() - def fsdp_kwargs(param_dtype, mesh): + def fsdp_kwargs(param_dtype): policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) return {"mp_policy": policy, "mesh": mesh} for unit in build_wrap_plan(model, compiled, list(model.blocks)): - fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype, shard_mesh if unit.shard else noshard_mesh)) - fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE, shard_mesh)) - - # Precision units are replicated, so their local shard is the whole tensor: no all-gather. - for unit in compiled.wrap_units: - param = next(unit.module.parameters()) - if not isinstance(param, DTensor) or param.placements[0] != Replicate(): - raise AssertionError(f"{unit.fqn} is not replicated: {getattr(param, 'placements', None)}") - if param.to_local().shape != param.shape: - raise AssertionError(f"{unit.fqn} local shard {param.to_local().shape} != full {param.shape}") + fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype)) + fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE)) seen: dict[str, torch.dtype] = {} diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index a91aa758..9b6bf483 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -62,7 +62,7 @@ def _units(compiled): def _plan(model, compiled): - return [(unit.fqn, unit.param_dtype, unit.shard) for unit in build_wrap_plan(model, compiled, list(model.blocks))] + return [(unit.fqn, unit.param_dtype) for unit in build_wrap_plan(model, compiled, list(model.blocks))] NORM_FQNS = {f"blocks.{i}{suffix}" for i in range(2) for suffix in (".norm", ".attn.norm_q")} @@ -192,7 +192,7 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): ├── attn [U] fp32 gather order 2 │ └── norm_q [U] bf16 gather order 1, wraps first └── rope fp16 buffer, never gathered - blocks.1 [U] bf16 block unit only, still sharded + blocks.1 [U] bf16 block unit only, at the default dtype """ model = _model() spec = PrecisionSpec( @@ -208,12 +208,11 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): "blocks.0.attn": torch.float32, "blocks.0": torch.float16, } - # (fqn, param_dtype, shard): precision units are replicated, the plain block unit shards. assert _plan(model, compiled) == [ - ("blocks.0.attn.norm_q", torch.bfloat16, False), - ("blocks.0.attn", torch.float32, False), - ("blocks.0", torch.float16, False), - ("blocks.1", torch.bfloat16, True), + ("blocks.0.attn.norm_q", torch.bfloat16), + ("blocks.0.attn", torch.float32), + ("blocks.0", torch.float16), + ("blocks.1", torch.bfloat16), ] @@ -232,9 +231,9 @@ def test_block_inside_an_override_wraps_at_the_override_dtype(): compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) assert _units(compiled) == [("blocks", torch.float32)] assert _plan(model, compiled) == [ - ("blocks.0", torch.float32, True), - ("blocks.1", torch.float32, True), - ("blocks", torch.float32, False), + ("blocks.0", torch.float32), + ("blocks.1", torch.float32), + ("blocks", torch.float32), ] From bd11d9d20bf092e2372202cb4c3c9c54d9e99fef Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 00:35:21 -0700 Subject: [PATCH 29/50] refactor(fsdp): reject an empty ModuleSel in __post_init__ and drop the rest of rule validation --- miles/backends/fsdp_utils/precision.py | 17 ++++------------- .../backends/fsdp_utils/test_precision_plan.py | 5 ++--- 2 files changed, 6 insertions(+), 16 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index be6295d9..b77fe384 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -75,6 +75,10 @@ class ModuleSel: fqn: str | None = None cls: str | None = None + def __post_init__(self) -> None: + if self.fqn is None and self.cls is None: + raise ValueError("ModuleSel needs fqn or cls; an empty selector silently matches every module") + @dataclass(frozen=True) class Rule: @@ -123,16 +127,6 @@ def apply_master_casts(self) -> None: cast.tensor.data = cast.tensor.data.to(cast.dtype) -def _validate_rule(rule: Rule) -> None: - if rule.master is None and rule.gather is None: - raise ValueError(f"precision rule sets no dtype axis: {rule}") - if rule.select.fqn is None and rule.select.cls is None: - raise ValueError(f"precision rule has an empty ModuleSel: {rule}") - for axis in (rule.master, rule.gather): - if axis is not None and axis != "default" and axis not in _DTYPES: - raise ValueError(f"precision rule has unknown dtype {axis!r}: {rule}") - - def _selects(sel: ModuleSel, mod_fqn: str, module: torch.nn.Module) -> bool: if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): return False @@ -175,9 +169,6 @@ def compile_precision_plan( default_dtype: torch.dtype, ) -> CompiledPrecision: """Resolve spec rules against the (pre-LoRA, pre-FSDP) model and lower them per module.""" - for rule in spec.rules: - _validate_rule(rule) - # (1) A rule matching nothing is almost certainly a typo'd pattern or class name. matched_fqns = [] for rule in spec.rules: diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index 9b6bf483..c814ac1d 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -173,9 +173,8 @@ def test_later_rule_overrides_earlier_selection(): def test_empty_module_sel_rejected(): """A selector with neither fqn nor cls would silently match every module.""" - spec = PrecisionSpec(rules=(Rule(ModuleSel(), master="fp32"),)) - with pytest.raises(ValueError, match="empty ModuleSel"): - compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + with pytest.raises(ValueError, match="needs fqn or cls"): + ModuleSel() def test_every_node_of_a_nested_chain_wraps_bottom_up(): From 9ce1785471df04bec296685031771fe01110e1f2 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 00:38:52 -0700 Subject: [PATCH 30/50] refactor(args): read the sglang dit precision directly instead of via getattr --- miles/utils/arguments.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 208e0279..5ba2bd39 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1510,8 +1510,7 @@ def set_default_diffusion_args(args) -> None: from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig # Mirrors the engine's forwarding rule: a value equal to the class default counts as unset. - sglang_dit = getattr(args, "sglang_dit_precision", None) - if sglang_dit is None or sglang_dit == getattr(PipelineConfig, "dit_precision", None): + if args.sglang_dit_precision == PipelineConfig.dit_precision: args.sglang_dit_precision = args.precision_default_dtype From 435f72e43063f1465024fc64f8e1ede39b76b9b5 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 00:43:57 -0700 Subject: [PATCH 31/50] fix(args): refuse a train/rollout dtype mismatch instead of overwriting sglang's per-pipeline default --- miles/utils/arguments.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 5ba2bd39..4ea19d8f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -196,9 +196,9 @@ def add_train_arguments(parser): choices=["fp16", "bf16", "fp32"], help=( "Default dtype for every module the family PrecisionSpec does not pin " - "explicitly: sets both the training-side forward/gather dtype " - "(--diffusion-forward-dtype) and the rollout engine's " - "--sglang-dit-precision. Specific flags win over this default." + "explicitly, i.e. the training-side forward/gather dtype when " + "--diffusion-forward-dtype is unset. Rollout must agree: pass the same value " + "to --sglang-dit-precision or startup is refused." ), ) parser.add_argument( @@ -1503,15 +1503,9 @@ def set_default_diffusion_args(args) -> None: else: args.ref_mode = "none" - # --precision-default-dtype fills every dtype knob left unset; specific flags win. + # --diffusion-forward-dtype wins over --precision-default-dtype, which wins over the built-in. if args.diffusion_forward_dtype is None: args.diffusion_forward_dtype = args.precision_default_dtype or "bf16" - if args.precision_default_dtype is not None: - from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig - - # Mirrors the engine's forwarding rule: a value equal to the class default counts as unset. - if args.sglang_dit_precision == PipelineConfig.dit_precision: - args.sglang_dit_precision = args.precision_default_dtype def miles_validate_args(args): @@ -1540,10 +1534,13 @@ def miles_validate_args(args): if args.eval_reward_key is None: args.eval_reward_key = args.reward_key - if args.precision_default_dtype is not None and args.sglang_dit_precision != args.precision_default_dtype: + # Rollout and training must denoise at the same dtype, or the PPO ratio compares log-probs from + # two different precisions. sglang resolves dit_precision from a per-pipeline config class we + # cannot see here, so refuse a mismatch rather than overwrite its value. + if args.sglang_dit_precision != args.diffusion_forward_dtype: raise ValueError( - f"--sglang-dit-precision {args.sglang_dit_precision} conflicts with " - f"--precision-default-dtype {args.precision_default_dtype}" + f"--sglang-dit-precision {args.sglang_dit_precision} disagrees with the training forward " + f"dtype {args.diffusion_forward_dtype}; pass both with the same value" ) args.update_weight_target_modules = [ From 1c8541f8e5f9cd423f8f83504a82f0fe0bb0aebb Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 00:45:59 -0700 Subject: [PATCH 32/50] chore(ltx): drop the redundant rollout dit precision flag --- miles/utils/arguments.py | 5 +++-- scripts/run-diffusion-grpo-ltx23-sglang.sh | 1 - 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4ea19d8f..075fce32 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1535,8 +1535,9 @@ def miles_validate_args(args): args.eval_reward_key = args.reward_key # Rollout and training must denoise at the same dtype, or the PPO ratio compares log-probs from - # two different precisions. sglang resolves dit_precision from a per-pipeline config class we - # cannot see here, so refuse a mismatch rather than overwrite its value. + # two different precisions. Recipes leave --sglang-dit-precision alone unless the model needs a + # non-default rollout dtype (SD3.5 is fp16), so this compares against the flag's default; a + # per-pipeline override inside sglang is resolved from the model path and is invisible here. if args.sglang_dit_precision != args.diffusion_forward_dtype: raise ValueError( f"--sglang-dit-precision {args.sglang_dit_precision} disagrees with the training forward " diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index f6693627..b1122861 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -82,7 +82,6 @@ fi --diffusion-forward-dtype bf16 \ --fsdp-master-dtype bf16 \ --fsdp-reduce-dtype bf16 \ - --sglang-dit-precision bf16 \ --advantage-estimator grpo \ --globalize-reward-std \ --rm-type pickscore \ From b2213de94c6906433c2bea8453eff95e86cd71a6 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 00:54:04 -0700 Subject: [PATCH 33/50] refactor(fsdp): resolve precision per module in two loops and hang wrap_plan off the result --- miles/backends/fsdp_utils/actor.py | 12 +- miles/backends/fsdp_utils/precision.py | 131 +++++++----------- .../fsdp_utils/_precision_wrap_worker.py | 6 +- .../fsdp_utils/test_precision_plan.py | 24 ++-- 4 files changed, 69 insertions(+), 104 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 77cc677b..6ad28781 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -39,13 +39,7 @@ from .lr_scheduler import get_lr_scheduler from .metrics import new_metric_buffer from .parallel import create_fsdp_parallel_state -from .precision import ( - apply_input_dtype_policy, - build_wrap_plan, - compile_precision_plan, - log_precision_summary, - resolve_dtype, -) +from .precision import apply_input_dtype_policy, compile_precision, log_precision_summary, resolve_dtype from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -136,7 +130,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.model_backend.enable_gradient_checkpointing(model) # Resolve the family precision plan on clean FQNs (pre-LoRA, pre-FSDP). - compiled_precision = compile_precision_plan( + compiled_precision = compile_precision( model, self.train_pipeline_config.precision_spec, default_dtype=self._forward_dtype, @@ -666,7 +660,7 @@ def _fsdp_kwargs(policy_param_dtype): "mesh": mesh, } - for unit in build_wrap_plan(model, compiled_precision, modules): + for unit in compiled_precision.wrap_plan(model, modules): fully_shard(unit.module, **_fsdp_kwargs(unit.param_dtype)) fully_shard(model, **_fsdp_kwargs(param_dtype)) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index b77fe384..d56aebcf 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -2,8 +2,7 @@ A family declares dtype intent as PrecisionSpec rules on its TrainPipelineConfig. A rule selects modules by FQN glob and/or class-name glob (both narrows to the -intersection) and pins one or both axes; a rule covers its modules' subtrees, and -where rules overlap the later one wins per axis: +intersection) and pins one or both axes: - master: resident dtype of the module's params/buffers (optimizer precision) - gather: dtype the params are cast to for FSDP all-gather + forward Compute dtype is not managed here: the trainer autocasts the DiT forward, @@ -11,24 +10,24 @@ ``apply_input_dtype_policy`` below, and op-level exceptions belong to the monkey-patch registry. -``compile_precision_plan`` lowers the rules onto what FSDP2 can express: +``compile_precision`` lowers the rules onto what FSDP2 can express: PrecisionSpec rules | - (1) match each rule to modules once - | - (2) per module, fold covering rules -> (master, gather) intent + (1) per module, parent-first: inherit the parent's decision, then apply the + rules selecting this module in spec order (so a deeper rule always wins + and order only breaks ties on one module) | | - master axis gather axis + master dtype gather dtype | | - (3) != loaded dtype? (4) != the dtype the enclosing wrap unit already - -> MasterCast of the provides? -> the module becomes its own wrap - module's own float unit at that dtype (paramless modules have - tensors, at load time nothing to gather and are skipped), which makes - | the units a minimal cover of the tree + (2) != loaded dtype? (3) != what the parent already provides? + -> MasterCast of the -> the module becomes its own wrap unit at that + module's own float dtype (paramless modules have nothing to gather + tensors, at load time and are skipped), which makes the units a + | minimal cover of the tree v | apply_master_casts() v - before FSDP wrapping build_wrap_plan() merges the units with the block + before FSDP wrapping compiled.wrap_plan() merges the units with the block modules into one deepest-first order, so FSDP2 always nests child-before-parent — that is how gather="default" carves a module back out of a @@ -126,6 +125,19 @@ def apply_master_casts(self) -> None: for cast in self.master_casts: cast.tensor.data = cast.tensor.data.to(cast.dtype) + def wrap_plan(self, model: torch.nn.Module, block_modules: list[torch.nn.Module]) -> list[WrapUnit]: + """One wrap order for FSDP2, deepest module first. Block modules are extra wraps that FSDP + needs for sharding granularity, so they must carry their own effective dtype: wrapping one at + the default inside an overridden region would be the innermost wrap and undo the override.""" + plan: dict[torch.nn.Module, WrapUnit] = {unit.module: unit for unit in self.wrap_units} + depths, fqns = {}, {} + for mod_fqn, module in model.named_modules(): + depths[module], fqns[module] = mod_fqn.count("."), mod_fqn + for module in block_modules: + fqn = fqns[module] + plan.setdefault(module, WrapUnit(fqn, module, self.gather_dtypes[fqn])) + return [plan[module] for module in sorted(plan, key=lambda module: -depths[module])] + def _selects(sel: ModuleSel, mod_fqn: str, module: torch.nn.Module) -> bool: if sel.fqn is not None and not fnmatch(mod_fqn, sel.fqn): @@ -138,93 +150,52 @@ def _parent_fqn(mod_fqn: str) -> str: return mod_fqn.rsplit(".", 1)[0] if "." in mod_fqn else "" -def _self_and_ancestors(mod_fqn: str) -> list[str]: - fqns = [] - while mod_fqn: - fqns.append(mod_fqn) - mod_fqn = _parent_fqn(mod_fqn) - return fqns + [""] - - -def _fold_covering_rules( - rules: tuple[Rule, ...], - matched_fqns: list[set[str]], - mod_fqn: str, -) -> tuple[str | None, str | None]: - """Later rules override earlier ones per axis; a rule covers its matched modules' subtrees.""" - covering = _self_and_ancestors(mod_fqn) - master = gather = None - for rule, fqns in zip(rules, matched_fqns, strict=True): - if fqns.isdisjoint(covering): - continue - master = rule.master if rule.master is not None else master - gather = rule.gather if rule.gather is not None else gather - return master, gather - - -def compile_precision_plan( +def compile_precision( model: torch.nn.Module, spec: PrecisionSpec, *, default_dtype: torch.dtype, ) -> CompiledPrecision: - """Resolve spec rules against the (pre-LoRA, pre-FSDP) model and lower them per module.""" - # (1) A rule matching nothing is almost certainly a typo'd pattern or class name. - matched_fqns = [] - for rule in spec.rules: - fqns = {fqn for fqn, module in model.named_modules() if _selects(rule.select, fqn, module)} - if not fqns: - raise ValueError(f"precision rule matched no module: {rule}") - matched_fqns.append(fqns) - - # (2)-(4) named_modules is parent-first, so the enclosing unit's dtype is already known here. + """Resolve the spec against the (pre-LoRA, pre-FSDP) model, then lower it per module.""" master_casts: list[MasterCast] = [] wrap_units: list[WrapUnit] = [] + masters: dict[str, torch.dtype | None] = {"": None} gather_dtypes: dict[str, torch.dtype] = {"": default_dtype} - for mod_fqn, module in model.named_modules(): - enclosing_dtype = gather_dtypes[_parent_fqn(mod_fqn)] - gather_dtypes[mod_fqn] = enclosing_dtype - master, gather = _fold_covering_rules(spec.rules, matched_fqns, mod_fqn) + hits = [0] * len(spec.rules) - master_dtype = _resolve_axis(master, default_dtype) - if master_dtype is not None: + for mod_fqn, module in model.named_modules(): + parent_fqn = _parent_fqn(mod_fqn) + master, gather = masters[parent_fqn], gather_dtypes[parent_fqn] + for i, rule in enumerate(spec.rules): + if not _selects(rule.select, mod_fqn, module): + continue + hits[i] += 1 + if rule.master is not None: + master = _resolve_axis(rule.master, default_dtype) + if rule.gather is not None: + gather = _resolve_axis(rule.gather, default_dtype) + masters[mod_fqn], gather_dtypes[mod_fqn] = master, gather_dtypes[parent_fqn] + + if master is not None: prefix = f"{mod_fqn}." if mod_fqn else "" own = list(module.named_parameters(recurse=False)) + list(module.named_buffers(recurse=False)) for name, tensor in own: - if tensor.is_floating_point() and tensor.dtype != master_dtype: - master_casts.append(MasterCast(f"{prefix}{name}", tensor, master_dtype)) + if tensor.is_floating_point() and tensor.dtype != master: + master_casts.append(MasterCast(f"{prefix}{name}", tensor, master)) - gather_dtype = _resolve_axis(gather, default_dtype) - if gather_dtype is None or gather_dtype == enclosing_dtype: - continue - if not any(param.is_floating_point() for param in module.parameters()): + if gather == gather_dtypes[parent_fqn] or not any(p.is_floating_point() for p in module.parameters()): continue if mod_fqn == "": raise ValueError("cannot wrap the root module for a gather override") - wrap_units.append(WrapUnit(mod_fqn, module, gather_dtype)) - gather_dtypes[mod_fqn] = gather_dtype + wrap_units.append(WrapUnit(mod_fqn, module, gather)) + gather_dtypes[mod_fqn] = gather + for rule, hit in zip(spec.rules, hits, strict=True): + if not hit: + raise ValueError(f"precision rule matched no module: {rule}") return CompiledPrecision(master_casts=master_casts, wrap_units=wrap_units, gather_dtypes=gather_dtypes) -def build_wrap_plan( - model: torch.nn.Module, - compiled: CompiledPrecision, - block_modules: list[torch.nn.Module], -) -> list[WrapUnit]: - """One wrap order for FSDP2, deepest module first. Block modules are extra wraps that FSDP - needs for sharding granularity, so they must carry their own effective dtype: wrapping one at - the default inside an overridden region would be the innermost wrap and undo the override.""" - plan: dict[torch.nn.Module, WrapUnit] = {unit.module: unit for unit in compiled.wrap_units} - depths, fqns = {}, {} - for mod_fqn, module in model.named_modules(): - depths[module], fqns[module] = mod_fqn.count("."), mod_fqn - for module in block_modules: - fqn = fqns[module] - plan.setdefault(module, WrapUnit(fqn, module, compiled.gather_dtypes[fqn])) - return [plan[module] for module in sorted(plan, key=lambda module: -depths[module])] - - def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: logger.info( f"precision[{component}]: default gather dtype {default_dtype}, " diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index d8d34a19..f9cf210f 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -32,7 +32,7 @@ from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard -from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, build_wrap_plan, compile_precision_plan +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision DEFAULT_DTYPE = torch.bfloat16 @@ -107,14 +107,14 @@ def main() -> None: mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("dp_shard",)) model = Net().to(torch.float32) # fp32 master - compiled = compile_precision_plan(model, SPEC, default_dtype=DEFAULT_DTYPE) + compiled = compile_precision(model, SPEC, default_dtype=DEFAULT_DTYPE) compiled.apply_master_casts() def fsdp_kwargs(param_dtype): policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) return {"mp_policy": policy, "mesh": mesh} - for unit in build_wrap_plan(model, compiled, list(model.blocks)): + for unit in compiled.wrap_plan(model, list(model.blocks)): fully_shard(unit.module, **fsdp_kwargs(unit.param_dtype)) fully_shard(model, **fsdp_kwargs(DEFAULT_DTYPE)) diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index c814ac1d..d80635e7 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -23,7 +23,7 @@ import torch import torch.nn as nn -from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, build_wrap_plan, compile_precision_plan +from miles.backends.fsdp_utils.precision import ModuleSel, PrecisionSpec, Rule, compile_precision class Rope(nn.Module): @@ -62,7 +62,7 @@ def _units(compiled): def _plan(model, compiled): - return [(unit.fqn, unit.param_dtype) for unit in build_wrap_plan(model, compiled, list(model.blocks))] + return [(unit.fqn, unit.param_dtype) for unit in compiled.wrap_plan(model, list(model.blocks))] NORM_FQNS = {f"blocks.{i}{suffix}" for i in range(2) for suffix in (".norm", ".attn.norm_q")} @@ -78,7 +78,7 @@ def test_empty_spec_compiles_to_nothing(): │ └── norm_q bf16 └── rope bf16 """ - compiled = compile_precision_plan(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) + compiled = compile_precision(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert compiled.wrap_units == [] @@ -97,7 +97,7 @@ def test_fqn_glob_selects_norms_across_depths(): """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), master="fp32", gather="fp32"),)) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) assert dict(_units(compiled)) == dict.fromkeys(NORM_FQNS, torch.float32) compiled.apply_master_casts() assert model.blocks[0].attn.norm_q.weight.dtype is torch.float32 @@ -117,7 +117,7 @@ def test_cls_glob_selects_by_class(): └── rope bf16 """ spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="*LayerNorm"), gather="fp32"),)) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + compiled = compile_precision(_model(), spec, default_dtype=torch.bfloat16) assert compiled.master_casts == [] assert {fqn for fqn, _ in _units(compiled)} == NORM_FQNS @@ -136,7 +136,7 @@ def test_master_rule_covers_matched_subtree(): """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) compiled.apply_master_casts() assert model.blocks[1].linear.weight.dtype is torch.float32 assert model.blocks[1].rope.freqs.dtype is torch.float32 @@ -162,7 +162,7 @@ def test_later_rule_overrides_earlier_selection(): Rule(ModuleSel(fqn="blocks.0.*norm*"), gather="fp16"), ) ) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + compiled = compile_precision(_model(), spec, default_dtype=torch.bfloat16) assert dict(_units(compiled)) == { "blocks.0.norm": torch.float16, "blocks.0.attn.norm_q": torch.float16, @@ -201,7 +201,7 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): Rule(ModuleSel(fqn="blocks.0.attn.norm_q"), gather="default"), ) ) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) assert dict(_units(compiled)) == { "blocks.0.attn.norm_q": torch.bfloat16, "blocks.0.attn": torch.float32, @@ -227,7 +227,7 @@ def test_block_inside_an_override_wraps_at_the_override_dtype(): """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks"), gather="fp32"),)) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) assert _units(compiled) == [("blocks", torch.float32)] assert _plan(model, compiled) == [ ("blocks.0", torch.float32), @@ -249,7 +249,7 @@ def test_inherited_gather_needs_no_extra_unit(): └── rope bf16 """ spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Attn"), gather="fp32"),)) - compiled = compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + compiled = compile_precision(_model(), spec, default_dtype=torch.bfloat16) assert _units(compiled) == [("blocks.0.attn", torch.float32), ("blocks.1.attn", torch.float32)] @@ -263,7 +263,7 @@ def test_buffer_only_module_casts_master_without_wrapping(): """ model = _model() spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32", gather="fp32"),)) - compiled = compile_precision_plan(model, spec, default_dtype=torch.bfloat16) + compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} assert compiled.wrap_units == [] @@ -272,4 +272,4 @@ def test_unmatched_rule_rejected(): """A rule matching nothing is a typo'd pattern or class name, not a silent no-op.""" spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) with pytest.raises(ValueError, match="matched no module"): - compile_precision_plan(_model(), spec, default_dtype=torch.bfloat16) + compile_precision(_model(), spec, default_dtype=torch.bfloat16) From ff9b6643d18c35b3ad68f3d82721dee9c1cf6a9f Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 00:59:24 -0700 Subject: [PATCH 34/50] docs(fsdp): explain each step inside compile_precision --- miles/backends/fsdp_utils/precision.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index d56aebcf..bf17a08b 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -156,7 +156,17 @@ def compile_precision( *, default_dtype: torch.dtype, ) -> CompiledPrecision: - """Resolve the spec against the (pre-LoRA, pre-FSDP) model, then lower it per module.""" + """Resolve the spec against the (pre-LoRA, pre-FSDP) model, then lower it per module. + + ``named_modules`` yields parents before children, so a module's parent is already resolved when + we reach it: inheriting the subtree decision is one dict lookup, and every rule only has to be + tested against the module it names. Two FQN-keyed dicts carry that state forward: + + masters the master dtype *intent*, inherited down the subtree so a rule on a container + also casts the tensors its descendants own + gather_dtypes the *effective* gather dtype, i.e. what the innermost enclosing wrap unit + provides — seeded from the parent, advanced only where a unit is emitted + """ master_casts: list[MasterCast] = [] wrap_units: list[WrapUnit] = [] masters: dict[str, torch.dtype | None] = {"": None} @@ -166,6 +176,8 @@ def compile_precision( for mod_fqn, module in model.named_modules(): parent_fqn = _parent_fqn(mod_fqn) master, gather = masters[parent_fqn], gather_dtypes[parent_fqn] + # Rules selecting this module apply in spec order, so the later one wins per axis; rules on + # ancestors already took effect through the inherited values above. for i, rule in enumerate(spec.rules): if not _selects(rule.select, mod_fqn, module): continue @@ -174,8 +186,10 @@ def compile_precision( master = _resolve_axis(rule.master, default_dtype) if rule.gather is not None: gather = _resolve_axis(rule.gather, default_dtype) + # Seed the effective dtype from the parent; it only advances if this module emits a unit. masters[mod_fqn], gather_dtypes[mod_fqn] = master, gather_dtypes[parent_fqn] + # Master is per tensor: cast what this module owns, descendants get their own turn. if master is not None: prefix = f"{mod_fqn}." if mod_fqn else "" own = list(module.named_parameters(recurse=False)) + list(module.named_buffers(recurse=False)) @@ -183,6 +197,8 @@ def compile_precision( if tensor.is_floating_point() and tensor.dtype != master: master_casts.append(MasterCast(f"{prefix}{name}", tensor, master)) + # Gather is per module: a unit is needed only where the dtype changes and there is something + # to gather at all, which is what keeps the units a minimal cover of the tree. if gather == gather_dtypes[parent_fqn] or not any(p.is_floating_point() for p in module.parameters()): continue if mod_fqn == "": @@ -190,6 +206,7 @@ def compile_precision( wrap_units.append(WrapUnit(mod_fqn, module, gather)) gather_dtypes[mod_fqn] = gather + # A rule that selected nothing is a typo'd pattern or class name, not a silent no-op. for rule, hit in zip(spec.rules, hits, strict=True): if not hit: raise ValueError(f"precision rule matched no module: {rule}") From 3248f19800949592076b67d90b8a313cfab515a8 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:34:16 -0700 Subject: [PATCH 35/50] refactor(fsdp): drop the per-module master axis, leaving gather as the only rule axis --- miles/backends/fsdp_utils/actor.py | 1 - miles/backends/fsdp_utils/precision.py | 113 ++++--------- .../fsdp_utils/_precision_wrap_worker.py | 3 +- .../fsdp_utils/test_precision_plan.py | 159 +++++++----------- 4 files changed, 100 insertions(+), 176 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 6ad28781..f9d691fa 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -135,7 +135,6 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty self.train_pipeline_config.precision_spec, default_dtype=self._forward_dtype, ) - compiled_precision.apply_master_casts() if rank == 0: log_precision_summary(component, compiled_precision, default_dtype=self._forward_dtype) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index bf17a08b..37dfdff2 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -1,11 +1,11 @@ """Fine-grained weight-precision control for FSDP2, at module granularity. -A family declares dtype intent as PrecisionSpec rules on its TrainPipelineConfig. -A rule selects modules by FQN glob and/or class-name glob (both narrows to the -intersection) and pins one or both axes: - - master: resident dtype of the module's params/buffers (optimizer precision) - - gather: dtype the params are cast to for FSDP all-gather + forward -Compute dtype is not managed here: the trainer autocasts the DiT forward, +A family declares gather dtypes as PrecisionSpec rules on its TrainPipelineConfig: +a rule selects modules by FQN glob and/or class-name glob (both narrows to the +intersection) and pins the dtype their params are cast to for the FSDP +all-gather and the forward. The resident (master) dtype is a run-level knob, +``--fsdp-master-dtype``, because no model wants it to vary per module; compute +dtype is not managed here either — the trainer autocasts the DiT forward, model-boundary input dtypes are family policy applied by ``apply_input_dtype_policy`` below, and op-level exceptions belong to the monkey-patch registry. @@ -14,28 +14,20 @@ PrecisionSpec rules | - (1) per module, parent-first: inherit the parent's decision, then apply the - rules selecting this module in spec order (so a deeper rule always wins - and order only breaks ties on one module) - | | - master dtype gather dtype - | | - (2) != loaded dtype? (3) != what the parent already provides? - -> MasterCast of the -> the module becomes its own wrap unit at that - module's own float dtype (paramless modules have nothing to gather - tensors, at load time and are skipped), which makes the units a - | minimal cover of the tree - v | - apply_master_casts() v - before FSDP wrapping compiled.wrap_plan() merges the units with the block - modules into one deepest-first order, so FSDP2 - always nests child-before-parent — that is how - gather="default" carves a module back out of a - non-default ancestor. - -Module granularity is the floor FSDP2 gives us (fully_shard wraps modules, and -FSDP2 requires uniform master dtype among trainable params per unit), so -finer-grained selectors are deliberately not offered. + (1) per module, parent-first: inherit the parent's gather dtype, then apply + the rules selecting this module in spec order (so a deeper rule always + wins and order only breaks ties on one module) + | + (2) differs from what the parent already provides? -> the module becomes its + own wrap unit at that dtype (paramless modules have nothing to gather and + are skipped), which makes the units a minimal cover of the tree + | + (3) compiled.wrap_plan() merges the units with the block modules into one + deepest-first order, so FSDP2 always nests child-before-parent — that is + how gather="default" carves a module back out of a non-default ancestor. + +Module granularity is the floor FSDP2 gives us: fully_shard wraps modules, so a +finer-grained selector could not be lowered. """ from __future__ import annotations @@ -81,11 +73,10 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class Rule: - """Axes take a dtype name ("fp32"/"bf16"/"fp16"), "default" (the run's default dtype), or None (untouched).""" + """gather is a dtype name ("fp32"/"bf16"/"fp16") or "default", the run's default dtype.""" select: ModuleSel - master: str | None = None - gather: str | None = None + gather: str @dataclass(frozen=True) @@ -94,17 +85,10 @@ class PrecisionSpec: # --------------------------------------------------------------------------- -# Compiler: spec -> FSDP2 lowering (master casts + per-module wrap units) +# Compiler: spec -> FSDP2 lowering (per-module wrap units) # --------------------------------------------------------------------------- -@dataclass(frozen=True) -class MasterCast: - fqn: str - tensor: torch.Tensor - dtype: torch.dtype - - @dataclass(frozen=True) class WrapUnit: """A module to fully_shard on its own with param_dtype=gather.""" @@ -116,15 +100,10 @@ class WrapUnit: @dataclass class CompiledPrecision: - master_casts: list[MasterCast] wrap_units: list[WrapUnit] # Effective gather dtype of every module, i.e. the dtype its innermost wrap unit provides. gather_dtypes: dict[str, torch.dtype] - def apply_master_casts(self) -> None: - for cast in self.master_casts: - cast.tensor.data = cast.tensor.data.to(cast.dtype) - def wrap_plan(self, model: torch.nn.Module, block_modules: list[torch.nn.Module]) -> list[WrapUnit]: """One wrap order for FSDP2, deepest module first. Block modules are extra wraps that FSDP needs for sharding granularity, so they must carry their own effective dtype: wrapping one at @@ -160,46 +139,30 @@ def compile_precision( ``named_modules`` yields parents before children, so a module's parent is already resolved when we reach it: inheriting the subtree decision is one dict lookup, and every rule only has to be - tested against the module it names. Two FQN-keyed dicts carry that state forward: - - masters the master dtype *intent*, inherited down the subtree so a rule on a container - also casts the tensors its descendants own - gather_dtypes the *effective* gather dtype, i.e. what the innermost enclosing wrap unit - provides — seeded from the parent, advanced only where a unit is emitted + tested against the module it names. ``gather_dtypes`` carries that state forward — the effective + dtype of each module, i.e. what its innermost enclosing wrap unit provides. """ - master_casts: list[MasterCast] = [] wrap_units: list[WrapUnit] = [] - masters: dict[str, torch.dtype | None] = {"": None} gather_dtypes: dict[str, torch.dtype] = {"": default_dtype} hits = [0] * len(spec.rules) for mod_fqn, module in model.named_modules(): parent_fqn = _parent_fqn(mod_fqn) - master, gather = masters[parent_fqn], gather_dtypes[parent_fqn] - # Rules selecting this module apply in spec order, so the later one wins per axis; rules on - # ancestors already took effect through the inherited values above. + inherited = gather_dtypes[parent_fqn] + # Rules selecting this module apply in spec order, so the later one wins; rules on ancestors + # already took effect through the inherited value. + gather = inherited for i, rule in enumerate(spec.rules): if not _selects(rule.select, mod_fqn, module): continue hits[i] += 1 - if rule.master is not None: - master = _resolve_axis(rule.master, default_dtype) - if rule.gather is not None: - gather = _resolve_axis(rule.gather, default_dtype) + gather = _resolve_axis(rule.gather, default_dtype) # Seed the effective dtype from the parent; it only advances if this module emits a unit. - masters[mod_fqn], gather_dtypes[mod_fqn] = master, gather_dtypes[parent_fqn] - - # Master is per tensor: cast what this module owns, descendants get their own turn. - if master is not None: - prefix = f"{mod_fqn}." if mod_fqn else "" - own = list(module.named_parameters(recurse=False)) + list(module.named_buffers(recurse=False)) - for name, tensor in own: - if tensor.is_floating_point() and tensor.dtype != master: - master_casts.append(MasterCast(f"{prefix}{name}", tensor, master)) - - # Gather is per module: a unit is needed only where the dtype changes and there is something - # to gather at all, which is what keeps the units a minimal cover of the tree. - if gather == gather_dtypes[parent_fqn] or not any(p.is_floating_point() for p in module.parameters()): + gather_dtypes[mod_fqn] = inherited + + # A unit is needed only where the dtype changes and there is something to gather at all, + # which is what keeps the units a minimal cover of the tree. + if gather == inherited or not any(p.is_floating_point() for p in module.parameters()): continue if mod_fqn == "": raise ValueError("cannot wrap the root module for a gather override") @@ -210,16 +173,14 @@ def compile_precision( for rule, hit in zip(spec.rules, hits, strict=True): if not hit: raise ValueError(f"precision rule matched no module: {rule}") - return CompiledPrecision(master_casts=master_casts, wrap_units=wrap_units, gather_dtypes=gather_dtypes) + return CompiledPrecision(wrap_units=wrap_units, gather_dtypes=gather_dtypes) def log_precision_summary(component: str, compiled: CompiledPrecision, *, default_dtype: torch.dtype) -> None: logger.info( f"precision[{component}]: default gather dtype {default_dtype}, " - f"{len(compiled.master_casts)} master casts, {len(compiled.wrap_units)} extra wrap units" + f"{len(compiled.wrap_units)} extra wrap units" ) - for cast in compiled.master_casts: - logger.info(f"precision[{component}]: master {cast.fqn} -> {cast.dtype}") for unit in compiled.wrap_units: logger.info(f"precision[{component}]: wrap {unit.fqn} @ {unit.param_dtype}") diff --git a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py index f9cf210f..c8111c5c 100644 --- a/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -1,6 +1,6 @@ """Gloo worker asserting the compiled precision plan really wraps under FSDP2 (2 ranks). -Master is fp32 everywhere, default gather dtype is bf16, and the spec is +Master is fp32 everywhere (--fsdp-master-dtype), default gather dtype is bf16, and the spec is Rule(cls="Norm", gather=fp32) Rule(fqn="blocks.0.attn", gather=fp16) @@ -108,7 +108,6 @@ def main() -> None: model = Net().to(torch.float32) # fp32 master compiled = compile_precision(model, SPEC, default_dtype=DEFAULT_DTYPE) - compiled.apply_master_casts() def fsdp_kwargs(param_dtype): policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=torch.float32, cast_forward_inputs=False) diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index d80635e7..c285c714 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -1,7 +1,7 @@ -"""Compiling PrecisionSpec rules into master casts and FSDP2 wrap units. +"""Compiling PrecisionSpec rules into FSDP2 wrap units. -Every test uses this model, loaded at bf16 with default_dtype=bf16, and every -docstring draws the resulting dtype per node (`[U]` = its own wrap unit). +Every test uses this model with default_dtype=bf16, and every docstring draws the +resulting gather dtype per node (`[U]` = the module becomes its own wrap unit). `blocks.1` mirrors `blocks.0`, so most diagrams only draw block 0. Tiny classes and own float tensors @@ -53,12 +53,8 @@ def __init__(self): self.blocks = nn.ModuleList([Block(), Block()]) -def _model(dtype=torch.bfloat16): - return Tiny().to(dtype) - - def _units(compiled): - return [(unit.fqn, unit.param_dtype) for unit in compiled.wrap_units] + return {unit.fqn: unit.param_dtype for unit in compiled.wrap_units} def _plan(model, compiled): @@ -71,37 +67,32 @@ def _plan(model, compiled): def test_empty_spec_compiles_to_nothing(): """No rules, so every node keeps the default and nothing is emitted. - blocks.0 gather bf16 + blocks.0 bf16 ├── linear bf16 ├── norm bf16 ├── attn bf16 │ └── norm_q bf16 └── rope bf16 """ - compiled = compile_precision(_model(), PrecisionSpec(), default_dtype=torch.bfloat16) - assert compiled.master_casts == [] + compiled = compile_precision(Tiny(), PrecisionSpec(), default_dtype=torch.bfloat16) assert compiled.wrap_units == [] def test_fqn_glob_selects_norms_across_depths(): """`*norm*` crosses dots, so one rule catches both norm depths and skips the Linear siblings. - Rule(fqn="*norm*", master=fp32, gather=fp32) + Rule(fqn="*norm*", gather=fp32) - blocks.0 master bf16 gather bf16 - ├── linear bf16 bf16 - ├── norm [U] fp32 fp32 - ├── attn bf16 bf16 - │ └── norm_q [U] fp32 fp32 - └── rope bf16 bf16 + blocks.0 bf16 + ├── linear bf16 + ├── norm [U] fp32 + ├── attn bf16 + │ └── norm_q [U] fp32 + └── rope bf16 """ - model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), master="fp32", gather="fp32"),)) - compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) - assert dict(_units(compiled)) == dict.fromkeys(NORM_FQNS, torch.float32) - compiled.apply_master_casts() - assert model.blocks[0].attn.norm_q.weight.dtype is torch.float32 - assert model.blocks[0].linear.weight.dtype is torch.bfloat16 + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="*norm*"), gather="fp32"),)) + compiled = compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) + assert _units(compiled) == dict.fromkeys(NORM_FQNS, torch.float32) def test_cls_glob_selects_by_class(): @@ -109,52 +100,45 @@ def test_cls_glob_selects_by_class(): Rule(cls="*LayerNorm", gather=fp32) - blocks.0 gather bf16 - ├── linear bf16 + blocks.0 bf16 ├── norm [U] fp32 - ├── attn bf16 - │ └── norm_q [U] fp32 - └── rope bf16 + └── attn + └── norm_q [U] fp32 """ spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="*LayerNorm"), gather="fp32"),)) - compiled = compile_precision(_model(), spec, default_dtype=torch.bfloat16) - assert compiled.master_casts == [] - assert {fqn for fqn, _ in _units(compiled)} == NORM_FQNS + compiled = compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) + assert set(_units(compiled)) == NORM_FQNS -def test_master_rule_covers_matched_subtree(): - """A rule covers its module's whole subtree, params and buffers alike, and stops at the sibling. +def test_rule_covers_the_matched_subtree(): + """A rule on a container hands its dtype to everything below, so one unit covers the subtree. - Rule(fqn="blocks.1", master=fp32) + Rule(fqn="blocks.1", gather=fp32) - blocks.0 master bf16 blocks.1 master fp32 - ├── linear bf16 ├── linear fp32 - ├── norm bf16 ├── norm fp32 - ├── attn bf16 ├── attn fp32 - │ └── norm_q bf16 │ └── norm_q fp32 - └── rope.freqs bf16 └── rope.freqs fp32 + blocks.0 bf16 blocks.1 [U] fp32 + ├── linear bf16 ├── linear fp32 (inherits, no unit) + ├── norm bf16 ├── norm fp32 (inherits, no unit) + └── attn bf16 └── attn fp32 (inherits, no unit) + └── norm_q bf16 └── norm_q fp32 (inherits, no unit) """ - model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), master="fp32"),)) + model = Tiny() + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), gather="fp32"),)) compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) - compiled.apply_master_casts() - assert model.blocks[1].linear.weight.dtype is torch.float32 - assert model.blocks[1].rope.freqs.dtype is torch.float32 - assert model.blocks[0].linear.weight.dtype is torch.bfloat16 + assert _units(compiled) == {"blocks.1": torch.float32} + assert compiled.gather_dtypes["blocks.1.attn.norm_q"] is torch.float32 + assert compiled.gather_dtypes["blocks.0.attn.norm_q"] is torch.bfloat16 def test_later_rule_overrides_earlier_selection(): - """Both rules cover block 0's norms; the later one wins there while block 1 keeps the first. + """Both rules select block 0's norms; the later one wins there while block 1 keeps the first. Rule(cls="LayerNorm", gather=fp32) # rule 1 Rule(fqn="blocks.0.*norm*", gather=fp16) # rule 2, wins where they overlap - blocks.0 gather bf16 blocks.1 gather bf16 - ├── linear bf16 ├── linear bf16 - ├── norm [U] fp16 (rule 2) ├── norm [U] fp32 (rule 1) - ├── attn bf16 ├── attn bf16 - │ └── norm_q [U] fp16 (rule 2) │ └── norm_q [U] fp32 (rule 1) - └── rope bf16 └── rope bf16 + blocks.0 blocks.1 + ├── norm [U] fp16 ├── norm [U] fp32 + └── attn └── attn + └── norm_q [U] fp16 └── norm_q [U] fp32 """ spec = PrecisionSpec( rules=( @@ -162,8 +146,8 @@ def test_later_rule_overrides_earlier_selection(): Rule(ModuleSel(fqn="blocks.0.*norm*"), gather="fp16"), ) ) - compiled = compile_precision(_model(), spec, default_dtype=torch.bfloat16) - assert dict(_units(compiled)) == { + compiled = compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) + assert _units(compiled) == { "blocks.0.norm": torch.float16, "blocks.0.attn.norm_q": torch.float16, "blocks.1.norm": torch.float32, @@ -185,15 +169,15 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): Rule(fqn="blocks.0.attn", gather=fp32) Rule(fqn="blocks.0.attn.norm_q", gather=default) # carved back out - blocks.0 [U] fp16 gather order 3 + blocks.0 [U] fp16 wrap order 3 ├── linear fp16 (inherits blocks.0) ├── norm fp16 (inherits blocks.0) - ├── attn [U] fp32 gather order 2 - │ └── norm_q [U] bf16 gather order 1, wraps first - └── rope fp16 buffer, never gathered + ├── attn [U] fp32 wrap order 2 + │ └── norm_q [U] bf16 wrap order 1, wraps first + └── rope buffer only, never gathered blocks.1 [U] bf16 block unit only, at the default dtype """ - model = _model() + model = Tiny() spec = PrecisionSpec( rules=( Rule(ModuleSel(fqn="blocks.0"), gather="fp16"), @@ -202,7 +186,7 @@ def test_every_node_of_a_nested_chain_wraps_bottom_up(): ) ) compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) - assert dict(_units(compiled)) == { + assert _units(compiled) == { "blocks.0.attn.norm_q": torch.bfloat16, "blocks.0.attn": torch.float32, "blocks.0": torch.float16, @@ -219,16 +203,16 @@ def test_block_inside_an_override_wraps_at_the_override_dtype(): """The rule sits above the block units, so the blocks wrap deeper than the override; at the default dtype they would be the innermost wrap and silently undo it. - Rule(fqn="blocks", gather=fp32) + Rule(fqn="blocks", gather=fp32) - blocks [U] fp32 gather order 3 (the override) - ├── 0 fp32 gather order 1, block unit forced to fp32 - └── 1 fp32 gather order 2, block unit forced to fp32 + blocks [U] fp32 wrap order 3 (the override) + ├── 0 fp32 wrap order 1, block unit forced to fp32 + └── 1 fp32 wrap order 2, block unit forced to fp32 """ - model = _model() + model = Tiny() spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks"), gather="fp32"),)) compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) - assert _units(compiled) == [("blocks", torch.float32)] + assert _units(compiled) == {"blocks": torch.float32} assert _plan(model, compiled) == [ ("blocks.0", torch.float32), ("blocks.1", torch.float32), @@ -236,40 +220,21 @@ def test_block_inside_an_override_wraps_at_the_override_dtype(): ] -def test_inherited_gather_needs_no_extra_unit(): - """Units are a minimal cover: a node whose dtype already comes from an enclosing unit is skipped. - - Rule(cls="Attn", gather=fp32) - - blocks.0 gather bf16 - ├── linear bf16 - ├── norm bf16 - ├── attn [U] fp32 (owns no param, but its subtree does) - │ └── norm_q fp32 inherits attn, so no unit of its own - └── rope bf16 - """ - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Attn"), gather="fp32"),)) - compiled = compile_precision(_model(), spec, default_dtype=torch.bfloat16) - assert _units(compiled) == [("blocks.0.attn", torch.float32), ("blocks.1.attn", torch.float32)] - - -def test_buffer_only_module_casts_master_without_wrapping(): - """Buffers are never gathered, so a paramless module takes the master cast and no unit. +def test_paramless_module_gets_no_unit(): + """Buffers are never gathered, so pinning a buffer-only module lowers to nothing. - Rule(cls="Rope", master=fp32, gather=fp32) + Rule(cls="Rope", gather=fp32) - blocks.0 master bf16 gather bf16 - └── rope.freqs fp32 n/a (buffer), no unit + blocks.0 + └── rope.freqs buffer -> no unit """ - model = _model() - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), master="fp32", gather="fp32"),)) - compiled = compile_precision(model, spec, default_dtype=torch.bfloat16) - assert {cast.fqn for cast in compiled.master_casts} == {"blocks.0.rope.freqs", "blocks.1.rope.freqs"} + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="Rope"), gather="fp32"),)) + compiled = compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) assert compiled.wrap_units == [] def test_unmatched_rule_rejected(): """A rule matching nothing is a typo'd pattern or class name, not a silent no-op.""" - spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), master="fp32"),)) + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), gather="fp32"),)) with pytest.raises(ValueError, match="matched no module"): - compile_precision(_model(), spec, default_dtype=torch.bfloat16) + compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) From cb027cd7873aaa6352a8dd67b24a7b26dabb7dad Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:41:57 -0700 Subject: [PATCH 36/50] docs(fsdp): split the wrap-unit skip conditions and name each reason --- miles/backends/fsdp_utils/precision.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 37dfdff2..39af54de 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -160,9 +160,12 @@ def compile_precision( # Seed the effective dtype from the parent; it only advances if this module emits a unit. gather_dtypes[mod_fqn] = inherited - # A unit is needed only where the dtype changes and there is something to gather at all, - # which is what keeps the units a minimal cover of the tree. - if gather == inherited or not any(p.is_floating_point() for p in module.parameters()): + # Minimal cover: a unit is only needed where the dtype changes. + if gather == inherited: + continue + # Nothing to gather: parameters() recurses so a container still counts, but buffers are + # never gathered and non-float params never cast, so wrapping those is pure overhead. + if not any(param.is_floating_point() for param in module.parameters()): continue if mod_fqn == "": raise ValueError("cannot wrap the root module for a gather override") From b1933ad3d71911cce80c39d3bd8d8341179ee4da Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:43:23 -0700 Subject: [PATCH 37/50] docs(fsdp): lead compile_precision with the differs-from-parent rule --- miles/backends/fsdp_utils/precision.py | 44 ++++++++++++-------------- 1 file changed, 21 insertions(+), 23 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 39af54de..295e08ad 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -135,42 +135,40 @@ def compile_precision( *, default_dtype: torch.dtype, ) -> CompiledPrecision: - """Resolve the spec against the (pre-LoRA, pre-FSDP) model, then lower it per module. + """Resolve the spec against the (pre-LoRA, pre-FSDP) model into FSDP2 wrap units. - ``named_modules`` yields parents before children, so a module's parent is already resolved when - we reach it: inheriting the subtree decision is one dict lookup, and every rule only has to be - tested against the module it names. ``gather_dtypes`` carries that state forward — the effective - dtype of each module, i.e. what its innermost enclosing wrap unit provides. + The rule is one line: **a module becomes its own wrap unit exactly when its gather dtype differs + from its parent's.** Anything matching its parent is already covered by the parent's unit, so the + emitted units are the minimal set of fully_shard calls that realises the spec. + + The traversal makes that cheap. ``named_modules`` yields parents before children, so the parent's + dtype is already in ``gather_dtypes`` when we reach a module: inheritance is one dict lookup, and + each rule only has to be tested against the module it names rather than against its ancestors. """ wrap_units: list[WrapUnit] = [] gather_dtypes: dict[str, torch.dtype] = {"": default_dtype} hits = [0] * len(spec.rules) for mod_fqn, module in model.named_modules(): - parent_fqn = _parent_fqn(mod_fqn) - inherited = gather_dtypes[parent_fqn] - # Rules selecting this module apply in spec order, so the later one wins; rules on ancestors - # already took effect through the inherited value. - gather = inherited + parent_gather = gather_dtypes[_parent_fqn(mod_fqn)] + # Start from what the parent provides, then let the rules selecting this module override it + # in spec order: the later one wins, and rules on ancestors already acted via the parent. + gather = parent_gather for i, rule in enumerate(spec.rules): if not _selects(rule.select, mod_fqn, module): continue hits[i] += 1 gather = _resolve_axis(rule.gather, default_dtype) - # Seed the effective dtype from the parent; it only advances if this module emits a unit. - gather_dtypes[mod_fqn] = inherited - - # Minimal cover: a unit is only needed where the dtype changes. - if gather == inherited: - continue - # Nothing to gather: parameters() recurses so a container still counts, but buffers are - # never gathered and non-float params never cast, so wrapping those is pure overhead. - if not any(param.is_floating_point() for param in module.parameters()): - continue - if mod_fqn == "": + + # Differs from the parent -> this module needs its own unit, unless it has nothing to gather: + # parameters() recurses so containers still count, but buffers are never gathered, so a + # buffer-only module like a RoPE cache would only add an empty FSDP group. + needs_unit = gather != parent_gather and next(module.parameters(), None) is not None + if needs_unit and mod_fqn == "": raise ValueError("cannot wrap the root module for a gather override") - wrap_units.append(WrapUnit(mod_fqn, module, gather)) - gather_dtypes[mod_fqn] = gather + if needs_unit: + wrap_units.append(WrapUnit(mod_fqn, module, gather)) + gather_dtypes[mod_fqn] = gather if needs_unit else parent_gather # A rule that selected nothing is a typo'd pattern or class name, not a silent no-op. for rule, hit in zip(spec.rules, hits, strict=True): From ce04b73977ae25e2466fceb58aca6028d6fcec60 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:44:25 -0700 Subject: [PATCH 38/50] docs(fsdp): move the compile details into the docstring and drop the inline comments --- miles/backends/fsdp_utils/precision.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 295e08ad..527696eb 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -144,6 +144,9 @@ def compile_precision( The traversal makes that cheap. ``named_modules`` yields parents before children, so the parent's dtype is already in ``gather_dtypes`` when we reach a module: inheritance is one dict lookup, and each rule only has to be tested against the module it names rather than against its ancestors. + Within a module the rules apply in spec order, so a later rule wins, while rules on ancestors + have already acted through the inherited dtype. A buffer-only module never needs a unit — FSDP + gathers parameters, not buffers — whereas a container does, since ``parameters()`` recurses. """ wrap_units: list[WrapUnit] = [] gather_dtypes: dict[str, torch.dtype] = {"": default_dtype} @@ -151,8 +154,6 @@ def compile_precision( for mod_fqn, module in model.named_modules(): parent_gather = gather_dtypes[_parent_fqn(mod_fqn)] - # Start from what the parent provides, then let the rules selecting this module override it - # in spec order: the later one wins, and rules on ancestors already acted via the parent. gather = parent_gather for i, rule in enumerate(spec.rules): if not _selects(rule.select, mod_fqn, module): @@ -160,9 +161,6 @@ def compile_precision( hits[i] += 1 gather = _resolve_axis(rule.gather, default_dtype) - # Differs from the parent -> this module needs its own unit, unless it has nothing to gather: - # parameters() recurses so containers still count, but buffers are never gathered, so a - # buffer-only module like a RoPE cache would only add an empty FSDP group. needs_unit = gather != parent_gather and next(module.parameters(), None) is not None if needs_unit and mod_fqn == "": raise ValueError("cannot wrap the root module for a gather override") From afd564ce2cbff2e844a3e2d56031e1926d094ad6 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:49:22 -0700 Subject: [PATCH 39/50] fix(args): only check the rollout dtype when it was explicitly chosen --- miles/utils/arguments.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 075fce32..7a38ca43 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1535,13 +1535,16 @@ def miles_validate_args(args): args.eval_reward_key = args.reward_key # Rollout and training must denoise at the same dtype, or the PPO ratio compares log-probs from - # two different precisions. Recipes leave --sglang-dit-precision alone unless the model needs a - # non-default rollout dtype (SD3.5 is fp16), so this compares against the flag's default; a - # per-pipeline override inside sglang is resolved from the model path and is invisible here. - if args.sglang_dit_precision != args.diffusion_forward_dtype: + # two different precisions. A value left at the parser default means "unspecified": sglang then + # resolves dit_precision from a per-pipeline config we cannot see here, so only an explicitly + # chosen rollout dtype is checked (this mirrors the engine's own forwarding rule). + from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig + + rollout_dtype = args.sglang_dit_precision + if rollout_dtype != PipelineConfig.dit_precision and rollout_dtype != args.diffusion_forward_dtype: raise ValueError( - f"--sglang-dit-precision {args.sglang_dit_precision} disagrees with the training forward " - f"dtype {args.diffusion_forward_dtype}; pass both with the same value" + f"--sglang-dit-precision {rollout_dtype} disagrees with the training forward dtype " + f"{args.diffusion_forward_dtype}; pass both with the same value" ) args.update_weight_target_modules = [ From 8a77ccdc5c6da8b9faef2279adab69bbdd1cfc4d Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:50:32 -0700 Subject: [PATCH 40/50] revert(args): drop the rollout dtype parity check --- miles/utils/arguments.py | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 7a38ca43..8bf25227 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -197,8 +197,8 @@ def add_train_arguments(parser): help=( "Default dtype for every module the family PrecisionSpec does not pin " "explicitly, i.e. the training-side forward/gather dtype when " - "--diffusion-forward-dtype is unset. Rollout must agree: pass the same value " - "to --sglang-dit-precision or startup is refused." + "--diffusion-forward-dtype is unset. The rollout engine has its own " + "--sglang-dit-precision; keep them consistent in the recipe." ), ) parser.add_argument( @@ -1534,19 +1534,6 @@ def miles_validate_args(args): if args.eval_reward_key is None: args.eval_reward_key = args.reward_key - # Rollout and training must denoise at the same dtype, or the PPO ratio compares log-probs from - # two different precisions. A value left at the parser default means "unspecified": sglang then - # resolves dit_precision from a per-pipeline config we cannot see here, so only an explicitly - # chosen rollout dtype is checked (this mirrors the engine's own forwarding rule). - from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig - - rollout_dtype = args.sglang_dit_precision - if rollout_dtype != PipelineConfig.dit_precision and rollout_dtype != args.diffusion_forward_dtype: - raise ValueError( - f"--sglang-dit-precision {rollout_dtype} disagrees with the training forward dtype " - f"{args.diffusion_forward_dtype}; pass both with the same value" - ) - args.update_weight_target_modules = [ name.strip() for name in args.update_weight_target_module.split(",") if name.strip() ] From 135e83fac6686f3b107bdecfb03b177c4171c38f Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:55:23 -0700 Subject: [PATCH 41/50] feat(args): make --precision-default-dtype set both sides and reject disagreement --- miles/utils/arguments.py | 25 ++++++++++++++++---- scripts/run-diffusion-grpo-sd3-ocr-sglang.sh | 3 +-- scripts/run-diffusion-nft-sd3-pickscore.sh | 3 +-- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 8bf25227..cc09b856 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -196,9 +196,9 @@ def add_train_arguments(parser): choices=["fp16", "bf16", "fp32"], help=( "Default dtype for every module the family PrecisionSpec does not pin " - "explicitly, i.e. the training-side forward/gather dtype when " - "--diffusion-forward-dtype is unset. The rollout engine has its own " - "--sglang-dit-precision; keep them consistent in the recipe." + "explicitly. One knob for both sides: it sets the training forward/gather dtype " + "and the rollout engine's --sglang-dit-precision. Leave it unset to tune each " + "side on its own, which is what the shipped recipes do." ), ) parser.add_argument( @@ -1407,6 +1407,8 @@ def add_sglang_tp_size(): ) parser.set_defaults(sglang_tensor_parallel_size=add_sglang_tp_size()) + # None means unset: the engine then forwards nothing and sglang keeps its per-pipeline dtype. + parser.set_defaults(sglang_dit_precision=None) return parser return add_miles_arguments @@ -1503,9 +1505,13 @@ def set_default_diffusion_args(args) -> None: else: args.ref_mode = "none" - # --diffusion-forward-dtype wins over --precision-default-dtype, which wins over the built-in. + # --precision-default-dtype fills whichever side was left unset; disagreements are rejected in + # miles_validate_args. Without it the training side falls back to bf16 and the rollout engine + # keeps its own per-pipeline default. if args.diffusion_forward_dtype is None: args.diffusion_forward_dtype = args.precision_default_dtype or "bf16" + if args.sglang_dit_precision is None: + args.sglang_dit_precision = args.precision_default_dtype def miles_validate_args(args): @@ -1534,6 +1540,17 @@ def miles_validate_args(args): if args.eval_reward_key is None: args.eval_reward_key = args.reward_key + if args.precision_default_dtype is not None: + for flag, value in ( + ("--diffusion-forward-dtype", args.diffusion_forward_dtype), + ("--sglang-dit-precision", args.sglang_dit_precision), + ): + if value != args.precision_default_dtype: + raise ValueError( + f"{flag} {value} disagrees with --precision-default-dtype " + f"{args.precision_default_dtype}; leave it unset or pass the same value" + ) + args.update_weight_target_modules = [ name.strip() for name in args.update_weight_target_module.split(",") if name.strip() ] diff --git a/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh b/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh index a98d0e13..e728eca5 100755 --- a/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh +++ b/scripts/run-diffusion-grpo-sd3-ocr-sglang.sh @@ -99,8 +99,7 @@ python -u "${ROOT_DIR}/train_diffusion.py" \ --advantage-estimator grpo \ --globalize-reward-std \ --rm-type ocr \ - --diffusion-forward-dtype fp16 \ - --sglang-dit-precision fp16 \ + --precision-default-dtype fp16 \ --sglang-vae-slicing \ --diffusion-num-steps 10 \ --diffusion-step-strategy-path miles.rollout.step_strategy_hub.sde_window \ diff --git a/scripts/run-diffusion-nft-sd3-pickscore.sh b/scripts/run-diffusion-nft-sd3-pickscore.sh index 947a071f..57392f33 100755 --- a/scripts/run-diffusion-nft-sd3-pickscore.sh +++ b/scripts/run-diffusion-nft-sd3-pickscore.sh @@ -128,8 +128,7 @@ python -u "${ROOT_DIR}/train_diffusion.py" \ --globalize-reward-std \ --diffusion-model "${SD3_MODEL}" \ "${REWARD_ARGS[@]}" \ - --diffusion-forward-dtype fp16 \ - --sglang-dit-precision fp16 \ + --precision-default-dtype fp16 \ --sglang-vae-slicing \ --diffusion-num-steps 10 \ --diffusion-eval-num-steps 50 \ From ceafe03b36a365dd6b654158adc4a7fe4128d86e Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 01:58:44 -0700 Subject: [PATCH 42/50] docs: cut the dtype cascade comment to one line and drop stale master wording --- miles/backends/fsdp_utils/actor.py | 2 +- miles/backends/fsdp_utils/configs/train_pipeline_config.py | 2 +- miles/utils/arguments.py | 4 +--- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f9d691fa..92d0a902 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -129,7 +129,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if args.gradient_checkpointing: self.model_backend.enable_gradient_checkpointing(model) - # Resolve the family precision plan on clean FQNs (pre-LoRA, pre-FSDP). + # Resolve the family precision spec on clean FQNs (pre-LoRA, pre-FSDP). compiled_precision = compile_precision( model, self.train_pipeline_config.precision_spec, diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index edace089..ccfee7d3 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -83,7 +83,7 @@ class TrainPipelineConfig(abc.ABC): supports_cfg_training: bool = True # Rollout parity patch group applied by the engine (see monkey_patches; None = none). rollout_patch_group: str | None = None - # Weight-precision rules (master/gather dtypes) compiled onto FSDP2; see precision.py. + # Gather-dtype rules compiled onto FSDP2 wrap units; see precision.py. precision_spec: PrecisionSpec = PrecisionSpec() # Model-boundary input dtypes (see precision.apply_input_dtype_policy); families opt into casts explicitly. input_dtype_policy: dict = {"latents": None, "cond": None, "timestep": None} diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index cc09b856..14d6df11 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1505,9 +1505,7 @@ def set_default_diffusion_args(args) -> None: else: args.ref_mode = "none" - # --precision-default-dtype fills whichever side was left unset; disagreements are rejected in - # miles_validate_args. Without it the training side falls back to bf16 and the rollout engine - # keeps its own per-pipeline default. + # --precision-default-dtype fills whichever side was left unset; validate rejects disagreements. if args.diffusion_forward_dtype is None: args.diffusion_forward_dtype = args.precision_default_dtype or "bf16" if args.sglang_dit_precision is None: From 11ebe4baa521df5c4fe19799bf8c4f9e1dc8dfce Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 02:09:14 -0700 Subject: [PATCH 43/50] revert(fsdp): keep the fsdp mesh dim names as they were --- miles/backends/fsdp_utils/parallel.py | 4 ++-- tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index f5fda80c..16a113e5 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -20,9 +20,9 @@ def build_fsdp_meshes( """Build the FSDP hybrid-shard and DP/SP views.""" world_mesh = init_device_mesh(device_type, (world_size,), mesh_dim_names=("world",)) - shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "dp_shard")) + shard_view = world_mesh._unflatten(0, (dp_replicate, world_size // dp_replicate), ("dp_replicate", "fsdp")) # A degree-1 replicate axis would all-reduce over a single rank every bucket. - fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["dp_shard"] + fsdp_mesh = shard_view if dp_replicate > 1 else shard_view["fsdp"] meshes = {"world": world_mesh, "fsdp": fsdp_mesh, "dp": world_mesh} if sp_size > 1: diff --git a/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py b/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py index da1dc15a..667af37a 100644 --- a/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py +++ b/tests/fast/backends/fsdp_utils/_hybrid_shard_mesh_worker.py @@ -61,11 +61,11 @@ def check(rank, world_size, dp_replicate, dp_shard, ring_degree, ulysses_degree) fsdp = meshes["fsdp"] if dp_replicate > 1: assert fsdp.ndim == 2, fsdp - assert fsdp.mesh_dim_names == ("dp_replicate", "dp_shard") + assert fsdp.mesh_dim_names == ("dp_replicate", "fsdp") assert fsdp["dp_replicate"].size() == dp_replicate - assert fsdp["dp_shard"].size() == dp_shard * sp_size + assert fsdp["fsdp"].size() == dp_shard * sp_size # dp_replicate outermost makes a shard group a contiguous rank run, i.e. one node. - shard_ranks = mesh_group_ranks(fsdp["dp_shard"]) + shard_ranks = mesh_group_ranks(fsdp["fsdp"]) assert shard_ranks == list(range(shard_ranks[0], shard_ranks[0] + len(shard_ranks))) else: assert fsdp.ndim == 1, fsdp From fdd65eef412a50657c2d92e2ee25d323b724f781 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 02:17:48 -0700 Subject: [PATCH 44/50] fix(fsdp): resolve wrap-plan block dtypes through the LoRA prefix --- miles/backends/fsdp_utils/precision.py | 3 ++- .../backends/fsdp_utils/test_precision_plan.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index 527696eb..a357753b 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -114,7 +114,8 @@ def wrap_plan(self, model: torch.nn.Module, block_modules: list[torch.nn.Module] depths[module], fqns[module] = mod_fqn.count("."), mod_fqn for module in block_modules: fqn = fqns[module] - plan.setdefault(module, WrapUnit(fqn, module, self.gather_dtypes[fqn])) + # The plan was compiled on the raw component; a later LoRA wrap prefixes the tree. + plan.setdefault(module, WrapUnit(fqn, module, self.gather_dtypes[fqn.removeprefix("base_model.model.")])) return [plan[module] for module in sorted(plan, key=lambda module: -depths[module])] diff --git a/tests/fast/backends/fsdp_utils/test_precision_plan.py b/tests/fast/backends/fsdp_utils/test_precision_plan.py index c285c714..3ff75bcf 100644 --- a/tests/fast/backends/fsdp_utils/test_precision_plan.py +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -238,3 +238,19 @@ def test_unmatched_rule_rejected(): spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="NoSuchModule"), gather="fp32"),)) with pytest.raises(ValueError, match="matched no module"): compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) + + +def test_wrap_plan_resolves_lora_prefixed_fqns(): + """The plan is compiled on the raw component; PEFT later nests it under + base_model.model, so wrap_plan must resolve block dtypes through the prefix.""" + model = Tiny() + compiled = compile_precision(model, PrecisionSpec(), default_dtype=torch.bfloat16) + inner = nn.Module() + inner.model = model + wrapped = nn.Module() + wrapped.base_model = inner + units = compiled.wrap_plan(wrapped, list(model.blocks)) + assert [(u.fqn, u.param_dtype) for u in units] == [ + ("base_model.model.blocks.0", torch.bfloat16), + ("base_model.model.blocks.1", torch.bfloat16), + ] From dc03e6b0e364b3ab59c249c2bc3de9bf5aedd8eb Mon Sep 17 00:00:00 2001 From: rockdu Date: Mon, 3 Aug 2026 17:53:15 -0700 Subject: [PATCH 45/50] feat(fsdp): pass rollout sigmas alongside timesteps to the DiT forward --- miles/backends/fsdp_utils/actor.py | 6 ++-- miles/backends/fsdp_utils/configs/ltx.py | 30 +++++++++++-------- .../backends/fsdp_utils/configs/qwen_image.py | 4 +++ miles/backends/fsdp_utils/configs/sd3.py | 1 - .../configs/train_pipeline_config.py | 5 ++-- miles/backends/fsdp_utils/configs/wan2_2.py | 1 - .../backends/fsdp_utils/loss_hub/flow_grpo.py | 16 ++++++---- miles/backends/fsdp_utils/loss_hub/nft.py | 10 ++----- miles/backends/fsdp_utils/loss_hub/types.py | 4 ++- miles/backends/fsdp_utils/precision.py | 8 +++-- .../test_train_pipeline_config_registry.py | 1 + .../fsdp_utils/test_flow_grpo_sigma_lookup.py | 28 +++++++++++++++++ .../fsdp_utils/test_input_dtype_policy.py | 25 +++++++++++----- 13 files changed, 97 insertions(+), 42 deletions(-) create mode 100644 tests/fast/backends/fsdp_utils/test_flow_grpo_sigma_lookup.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 92d0a902..36386b64 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -524,10 +524,11 @@ def _forward_train_pair_batch( forward_dtype = self._forward_dtype # Boundary dtypes are family policy; op interiors stay autocast-managed. - latents_in, timesteps_in, (pos_cond_in, neg_cond_in, joint_cond_in) = apply_input_dtype_policy( + latents_in, timesteps_in, sigmas_in, (pos_cond_in, neg_cond_in, joint_cond_in) = apply_input_dtype_policy( train_pipeline_config.input_dtype_policy, latents=prepared.latents, - timesteps=prepared.timesteps_for_model, + timesteps=prepared.timesteps, + sigmas=prepared.sigmas, conds=(prepared.pos_cond, prepared.neg_cond, prepared.joint_cond), default_dtype=forward_dtype, ) @@ -538,6 +539,7 @@ def _compute_noise_pred() -> torch.Tensor: model=prepared.model, latents_input=latents_in, timesteps_input=timesteps_in, + sigmas_input=sigmas_in, pos_cond=pos_cond_in, neg_cond=neg_cond_in, joint_cond=joint_cond_in, diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index f4a2604c..67b12f6e 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/miles/backends/fsdp_utils/configs/ltx.py @@ -13,11 +13,17 @@ @register_train_pipeline_config("ltx") class LTXTrainPipelineConfig(TrainPipelineConfig): - """LTX-2.3 video GRPO: unguided velocity forward over ltx_core.""" + """LTX-2.3 video GRPO: unguided velocity forward over ltx_core. + + Dtype parity vs sglang-d rollout (dump-verified on paired LTX-2.3 runs): the empty + precision_spec matches, and of the boundary axes only latents is load-bearing -- + forward_velocity anchors the element-wise math on latents.dtype and consumes + sigmas_input verbatim in fp32. Known benign delta: norm_out records fp32 under + autocast, compute-equivalent. + """ - needs_timestep_scaling = False supports_cfg_training = False - # Rollout stores σ×1000 in trajectory timesteps; ltx_core AdaLN uses σ∈[0,1]. + # Rollout stores σ×1000 in trajectory timesteps; the CPS SDE path resolves σ linearly. sde_timestep_divisor = 1000.0 rollout_patch_group = "ltx" hf_ckpt_name_patterns = ("ltx",) @@ -93,6 +99,7 @@ def compute_noise_pred( model: torch.nn.Module, latents_input: torch.Tensor, timesteps_input: torch.Tensor, + sigmas_input: torch.Tensor, pos_cond: dict | None, neg_cond: dict | None, joint_cond: dict | None, @@ -119,13 +126,13 @@ def compute_noise_pred( device=latents_input.device, dtype=latents_input.dtype, ) - return self.forward_velocity(model, latents_input, timesteps_input, cond) + return self.forward_velocity(model, latents_input, sigmas_input, cond) def forward_velocity( self, model: torch.nn.Module, latents_input: torch.Tensor, - timesteps_input: torch.Tensor, + sigmas_input: torch.Tensor, cond: dict, ) -> torch.Tensor: from ltx_core.model.transformer.modality import Modality @@ -134,17 +141,16 @@ def forward_velocity( dtype = latents_input.dtype B = latents_input.shape[0] - # Trajectory timesteps are σ×1000; ltx_core AdaLN expects σ∈[0,1] and - # multiplies by timestep_scale_multiplier (1000) internally. - sigma_scaled = timesteps_input.to(latents_input.dtype) - sigma_unit = sigma_scaled / float(self.sde_timestep_divisor) - per_token_t = sigma_unit.view(B, 1).to(dtype) + # The model consumes the rollout σ verbatim in fp32: bf16-rounding it before + # the sinusoid costs ~2e-3 rel per AdaLN block. + sigma_unit = sigmas_input.to(dtype) + per_token_t = sigma_unit.view(B, 1) video_modality = Modality( enabled=True, latent=latents_input, - sigma=sigma_unit.reshape(B), - timesteps=per_token_t, + sigma=sigmas_input.float().reshape(B), + timesteps=sigmas_input.float().view(B, 1), positions=cond["positions"].to(dtype), context=cond["context"].to(dtype), context_mask=None, diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 42910804..71310526 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -58,6 +58,10 @@ def _params(index: torch.Tensor, dim: int, theta: float = theta) -> torch.Tensor class QwenImageTrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("qwen-image",) + def compute_noise_pred(self, *, timesteps_input, sigmas_input, **kwargs): + # The QwenImage DiT consumes sigma as its timestep input. + return super().compute_noise_pred(timesteps_input=sigmas_input, sigmas_input=sigmas_input, **kwargs) + lora_target_modules = [ "to_q", "to_k", diff --git a/miles/backends/fsdp_utils/configs/sd3.py b/miles/backends/fsdp_utils/configs/sd3.py index 332e9fa1..cd35966d 100644 --- a/miles/backends/fsdp_utils/configs/sd3.py +++ b/miles/backends/fsdp_utils/configs/sd3.py @@ -25,7 +25,6 @@ class SD3TrainPipelineConfig(TrainPipelineConfig): "attn.add_v_proj", "attn.to_add_out", ] - needs_timestep_scaling = False def prepare_cond_kwargs(self, cond: CondKwargs | None, device: torch.device) -> dict: if cond is None: diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index ccfee7d3..980a50e8 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -76,7 +76,6 @@ class TrainPipelineConfig(abc.ABC): """Base class. Subclass per model family.""" lora_target_modules: list[str] = ["to_q", "to_k", "to_v", "to_out.0"] - needs_timestep_scaling: bool = True optimizer_state_allowed_missing: list[str] = [] # Case-insensitive substrings matched against the checkpoint name (--diffusion-model). hf_ckpt_name_patterns: tuple[str, ...] = () @@ -107,6 +106,7 @@ def compute_noise_pred( model: torch.nn.Module, latents_input: torch.Tensor, timesteps_input: torch.Tensor, + sigmas_input: torch.Tensor, pos_cond: dict | None, neg_cond: dict | None, joint_cond: dict | None, @@ -115,7 +115,8 @@ def compute_noise_pred( guidance_scale: float, true_cfg_scale: float | None, ) -> torch.Tensor: - """Default diffusers forward with CFG; families with a different forward override.""" + """Default diffusers forward with CFG; families whose model consumes sigma + (or with a different forward entirely) override.""" def _forward(cond: dict) -> torch.Tensor: return model( diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 7e227cd2..fa306167 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -15,7 +15,6 @@ class Wan2_2TrainPipelineConfig(TrainPipelineConfig): # ("transformer_2") the rest. boundary_ratio = 0.875 # Wan DiT expects raw scheduler timesteps (0..num_train_timesteps), no /1000 scaling. - needs_timestep_scaling = False def component_for_timestep(self, timestep: float, num_train_timesteps: int) -> str: if timestep >= self.boundary_ratio * num_train_timesteps: diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 876d69dc..3538d676 100644 --- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py +++ b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py @@ -14,6 +14,15 @@ def _stack_pair_field(batch: list[dict], key: str, device: torch.device) -> torc return torch.stack([pair[key] for pair in batch]).to(device=device, dtype=torch.float32) +def _sigmas_for_timesteps(scheduler, timesteps: torch.Tensor) -> torch.Tensor: + """Exact-match pair timesteps into the rollout scheduler snapshot and return their sigmas.""" + sched_t = scheduler.timesteps.to(timesteps.device) + idx = (timesteps.view(-1, 1) == sched_t.view(1, -1)).long().argmax(dim=1) + if not torch.equal(sched_t[idx], timesteps): + raise ValueError("train pair timesteps not found in rollout scheduler_timesteps") + return scheduler.sigmas.to(timesteps.device)[idx] + + def prepare_flow_grpo_batch( ctx: DiffusionLossContext, batch: list[dict], @@ -62,10 +71,7 @@ def prepare_flow_grpo_batch( args.diffusion_guidance_scale_2, ) - if config.needs_timestep_scaling: - timesteps_for_model = timesteps / float(num_train_timesteps) - else: - timesteps_for_model = timesteps + sigmas = _sigmas_for_timesteps(ctx.scheduler, timesteps) pos_list = [config.prepare_cond_kwargs(batch[i]["denoising_env"].pos_cond_kwargs, device) for i in range(bsz)] neg_list = ( @@ -86,7 +92,7 @@ def prepare_flow_grpo_batch( return PreparedBatch( latents=latents, timesteps=timesteps, - timesteps_for_model=timesteps_for_model, + sigmas=sigmas, model=model, component_name=component_name, guidance_scale=guidance_scale, diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 97932078..8243a36e 100644 --- a/miles/backends/fsdp_utils/loss_hub/nft.py +++ b/miles/backends/fsdp_utils/loss_hub/nft.py @@ -41,16 +41,12 @@ def prepare_nft_batch( pos_cond = config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len) num_train_timesteps = ctx.scheduler.config.num_train_timesteps - if config.needs_timestep_scaling: - timesteps_for_model = t.to(dtype=torch.float32) - else: - timesteps_for_model = t * float(num_train_timesteps) xt = corrupt(x0, t, sample_noise(x0)) return PreparedBatch( latents=xt, - timesteps=t, - timesteps_for_model=timesteps_for_model, + timesteps=t * float(num_train_timesteps), + sigmas=t, model=model, component_name=component_name, guidance_scale=0.0, @@ -125,7 +121,7 @@ def nft_loss_formula( use_adaptive = args.diffusion_nft_adaptive_weight x0 = prepared.extras["x0"] - t = prepared.timesteps + t = prepared.sigmas t_exp = t.view(len(batch), *([1] * (x0.ndim - 1))) r = nft_r_from_advantages(prepared.advantage, adv_clip_max=adv_clip_max) pos_loss, neg_loss = nft_branch_losses( diff --git a/miles/backends/fsdp_utils/loss_hub/types.py b/miles/backends/fsdp_utils/loss_hub/types.py index 33d4f767..5600b9b8 100644 --- a/miles/backends/fsdp_utils/loss_hub/types.py +++ b/miles/backends/fsdp_utils/loss_hub/types.py @@ -28,8 +28,10 @@ class PreparedBatch: """Actor-owned DiT forward inputs produced by a prepare hook.""" latents: torch.Tensor + # Per-pair rollout timesteps (t domain) and scheduler sigmas (sigma domain); each + # family's compute_noise_pred consumes whichever domain its model expects. timesteps: torch.Tensor - timesteps_for_model: torch.Tensor + sigmas: torch.Tensor model: nn.Module component_name: str guidance_scale: float diff --git a/miles/backends/fsdp_utils/precision.py b/miles/backends/fsdp_utils/precision.py index a357753b..4bd20fcf 100644 --- a/miles/backends/fsdp_utils/precision.py +++ b/miles/backends/fsdp_utils/precision.py @@ -197,11 +197,13 @@ def apply_input_dtype_policy( *, latents: torch.Tensor, timesteps: torch.Tensor, + sigmas: torch.Tensor, conds: tuple, default_dtype: torch.dtype, -) -> tuple[torch.Tensor, torch.Tensor, tuple]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple]: """Cast float boundary inputs per family policy ("default"/dtype name/None=passthrough); - autocast alone would leave element-wise ops running at the raw input dtype.""" + autocast alone would leave element-wise ops running at the raw input dtype. The "timestep" + axis covers both the timesteps and the rollout sigmas.""" unknown = set(policy) - set(INPUT_DTYPE_POLICY_KEYS) if unknown: raise ValueError(f"input_dtype_policy has unknown keys {sorted(unknown)}; known: {INPUT_DTYPE_POLICY_KEYS}") @@ -221,4 +223,4 @@ def _cast(value, dtype: torch.dtype | None): out_conds = tuple( None if cond is None else {key: _cast(value, cond_dtype) for key, value in cond.items()} for cond in conds ) - return _cast(latents, latents_dtype), _cast(timesteps, timestep_dtype), out_conds + return _cast(latents, latents_dtype), _cast(timesteps, timestep_dtype), _cast(sigmas, timestep_dtype), out_conds diff --git a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py index 15f11d4f..69cd87ef 100644 --- a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py +++ b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py @@ -64,6 +64,7 @@ def _call(self, **overrides): model=_CondBiasDiT(), latents_input=self.h, timesteps_input=torch.tensor([3.0, 5.0]), + sigmas_input=torch.tensor([0.003, 0.005]), pos_cond=self.pos, neg_cond=self.neg, joint_cond=None, diff --git a/tests/fast/backends/fsdp_utils/test_flow_grpo_sigma_lookup.py b/tests/fast/backends/fsdp_utils/test_flow_grpo_sigma_lookup.py new file mode 100644 index 00000000..7ce3cfd5 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_flow_grpo_sigma_lookup.py @@ -0,0 +1,28 @@ +from types import SimpleNamespace + +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[]) + +import pytest +import torch + +from miles.backends.fsdp_utils.loss_hub.flow_grpo import _sigmas_for_timesteps + + +def _scheduler(): + sigmas = torch.tensor([0.9782581925392151, 0.8, 0.5, 0.0]) + return SimpleNamespace(timesteps=sigmas[:-1] * 1000, sigmas=sigmas) + + +def test_exact_sigma_lookup(): + scheduler = _scheduler() + timesteps = torch.stack([scheduler.timesteps[2], scheduler.timesteps[0]]) + sigmas = _sigmas_for_timesteps(scheduler, timesteps) + assert torch.equal(sigmas, torch.stack([scheduler.sigmas[2], scheduler.sigmas[0]])) + + +def test_unmatched_timestep_rejected(): + scheduler = _scheduler() + with pytest.raises(ValueError, match="not found"): + _sigmas_for_timesteps(scheduler, torch.tensor([123.456])) diff --git a/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py index facd63de..f47cb30a 100644 --- a/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py +++ b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py @@ -14,48 +14,56 @@ def _inputs(): latents = torch.zeros(1, 4, 8, dtype=torch.float32) timesteps = torch.tensor([857.69], dtype=torch.float32) + sigmas = torch.tensor([0.85769], dtype=torch.float32) pos_cond = { "context": torch.zeros(1, 3, 8, dtype=torch.float32), "context_mask": torch.ones(1, 3, dtype=torch.int64), } - return latents, timesteps, pos_cond + return latents, timesteps, sigmas, pos_cond def test_default_policy_is_passthrough(): - latents, timesteps, pos_cond = _inputs() - out_latents, out_timesteps, (out_pos, out_neg, out_joint) = apply_input_dtype_policy( + """The base default pins nothing, so every float input keeps the dtype it arrived with.""" + latents, timesteps, sigmas, pos_cond = _inputs() + out_latents, out_timesteps, out_sigmas, (out_pos, out_neg, out_joint) = apply_input_dtype_policy( DEFAULT_POLICY, latents=latents, timesteps=timesteps, + sigmas=sigmas, conds=(pos_cond, None, None), default_dtype=torch.bfloat16, ) assert out_latents.dtype == torch.float32 assert out_timesteps.dtype == torch.float32 + assert out_sigmas.dtype == torch.float32 assert out_pos["context"].dtype == torch.float32 assert out_pos["context_mask"].dtype == torch.int64 assert out_neg is None and out_joint is None def test_family_override_timestep_default(): - latents, timesteps, pos_cond = _inputs() + latents, timesteps, sigmas, pos_cond = _inputs() policy = {**DEFAULT_POLICY, "timestep": "default"} - _, out_timesteps, _ = apply_input_dtype_policy( + _, out_timesteps, out_sigmas, _ = apply_input_dtype_policy( policy, latents=latents, timesteps=timesteps, + sigmas=sigmas, conds=(pos_cond, None, None), default_dtype=torch.bfloat16, ) assert out_timesteps.dtype == torch.bfloat16 + assert out_sigmas.dtype == torch.bfloat16 def test_cast_policy_casts_floats_only(): - latents, timesteps, pos_cond = _inputs() - out_latents, _, (out_pos, _, _) = apply_input_dtype_policy( + """A family that pins latents/cond gets them cast to the forward dtype; masks stay int.""" + latents, timesteps, sigmas, pos_cond = _inputs() + out_latents, _, _, (out_pos, _, _) = apply_input_dtype_policy( {"latents": "default", "cond": "default", "timestep": "fp32"}, latents=latents, timesteps=timesteps, + sigmas=sigmas, conds=(pos_cond, None, None), default_dtype=torch.bfloat16, ) @@ -65,12 +73,13 @@ def test_cast_policy_casts_floats_only(): def test_unknown_key_rejected(): - latents, timesteps, pos_cond = _inputs() + latents, timesteps, sigmas, pos_cond = _inputs() with pytest.raises(ValueError, match="unknown keys"): apply_input_dtype_policy( {"latnets": "default"}, latents=latents, timesteps=timesteps, + sigmas=sigmas, conds=(pos_cond, None, None), default_dtype=torch.bfloat16, ) From 64d14300b551d673f858b2da7834a61aab48ff0e Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 02:08:54 -0700 Subject: [PATCH 46/50] feat(ltx): train the LTX recipe with fp32 master weights --- scripts/run-diffusion-grpo-ltx23-sglang.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run-diffusion-grpo-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index b1122861..ecbca076 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -80,7 +80,7 @@ fi --diffusion-output-num-frames 57 \ --diffusion-fps 24 \ --diffusion-forward-dtype bf16 \ - --fsdp-master-dtype bf16 \ + --fsdp-master-dtype fp32 \ --fsdp-reduce-dtype bf16 \ --advantage-estimator grpo \ --globalize-reward-std \ From 4666f445492a7f61bf7956b19dc2df3bcc4ea408 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 02:35:41 -0700 Subject: [PATCH 47/50] test(e2e): re-record the LTX-2.3 standard for fp32 master weights --- .../test_ltx23_pickscore_grpo_4xGPU.json | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json b/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json index 8e720e9c..c97230ee 100644 --- a/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json +++ b/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json @@ -1,6 +1,6 @@ { "meta": { - "commit": "5a1b7986af06e3d80368935c88acdd3dc6d98a40", + "commit": "64d14300b551d673f858b2da7834a61aab48ff0e", "source": "test_ltx23_pickscore_grpo_4xGPU.py" }, "metrics": { @@ -11,7 +11,7 @@ ], [ 1, - 0.6819055080413818 + 0.6817942261695862 ] ], "rollout/reward/raw_median": [ @@ -21,7 +21,7 @@ ], [ 1, - 0.6862333416938782 + 0.6878968477249146 ] ], "rollout/reward/raw_num_samples": [ @@ -41,61 +41,61 @@ ], [ 1, - 0.042159032076597214 + 0.041681572794914246 ] ], "train/grad_norm": [ [ 1.0, - 1.147376860899385e-05 + 2.8766138711944222e-05 ], [ 2.0, - 1.245457497134339e-05 + 1.2603897630469874e-05 ], [ 3.0, - 2.025627691182308e-05 + 1.8178370737587102e-05 ], [ 4.0, - 1.9314091332489625e-05 + 2.2844391423859634e-05 ] ], "train/log_prob_mean_abs_diff": [ [ 1.0, - 1.0107954343159993e-06 + 3.625949223836263e-07 ], [ 2.0, - 1.188988486925761e-06 + 7.146348555882772e-07 ], [ 3.0, - 1.8222878376642864e-06 + 1.1622905731201172e-06 ], [ 4.0, - 1.761441429456075e-06 + 1.274670163790385e-06 ] ], "train/log_prob_new_idx_0": [ [ 1.0, - -0.8474982244273027 + -0.8474981039762497 ], [ 2.0, - -0.8480804494271675 + -0.8480802631626526 ], [ 3.0, - -0.7732201963663101 + -0.7732198803375164 ], [ 4.0, - -0.7731628809124231 + -0.7731634061783552 ] ], "train/log_prob_old_idx_0": [ From beb9ce4f917cc292b45d0e777df9df98358b8cad Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 11:12:13 -0700 Subject: [PATCH 48/50] fix(ltx): drop all-true attention masks in rollout for train kernel parity --- .../monkey_patches/__init__.py | 2 ++ .../patch_ltx2_trivial_attention_mask.py | 34 +++++++++++++++++++ 2 files changed, 36 insertions(+) create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_trivial_attention_mask.py diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 2755e204..d11d6b80 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -60,10 +60,12 @@ def apply_ltx2_rollout_patches() -> None: from miles.backends.sglang_diffusion_utils.monkey_patches import ( patch_ltx2_disable_av_cross, patch_ltx2_rollout_cond_kwargs, + patch_ltx2_trivial_attention_mask, ) patch_ltx2_rollout_cond_kwargs.apply() patch_ltx2_disable_av_cross.apply() + patch_ltx2_trivial_attention_mask.apply() def apply_env_selected_rollout_patches() -> None: diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_trivial_attention_mask.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_trivial_attention_mask.py new file mode 100644 index 00000000..e4e80819 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_trivial_attention_mask.py @@ -0,0 +1,34 @@ +"""Align LTX2Attention's SDPA dispatch with training. + +Training runs mask-free flash-family attention everywhere (context_mask=None). +Rollout's cross-attention keeps an all-ones mask (flash-ineligible) and then +follows torch's default SDPA priority, which ranks cuDNN first on Hopper — +a different kernel worth ~1e-3 rel per block. Drop trivial masks and pin the +flash-first priority for the maskless path. +""" + +from __future__ import annotations + + +def apply() -> None: + from torch.nn.attention import SDPBackend, sdpa_kernel + + from sglang.multimodal_gen.runtime.models.dits import ltx_2 + + orig_forward = ltx_2.LTX2Attention.forward + flash_first = [ + SDPBackend.FLASH_ATTENTION, + SDPBackend.CUDNN_ATTENTION, + SDPBackend.EFFICIENT_ATTENTION, + SDPBackend.MATH, + ] + + def forward(self, x, context=None, mask=None, **kwargs): + if mask is not None and bool(mask.all()): + mask = None + if mask is None: + with sdpa_kernel(flash_first, set_priority=True): + return orig_forward(self, x, context=context, mask=mask, **kwargs) + return orig_forward(self, x, context=context, mask=mask, **kwargs) + + ltx_2.LTX2Attention.forward = forward From b940e5ac5852f54d3e6ca7d709b8fe0e2cffef37 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 13:20:24 -0700 Subject: [PATCH 49/50] fix(ltx): run rollout output layernorms in fp32 to match autocast --- .../monkey_patches/__init__.py | 2 + .../patch_ltx2_norm_out_fp32.py | 37 +++++++++++++++++++ 2 files changed, 39 insertions(+) create mode 100644 miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_norm_out_fp32.py diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index d11d6b80..170ea7b4 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -59,6 +59,7 @@ def apply_sgld_monkey_patches() -> None: def apply_ltx2_rollout_patches() -> None: from miles.backends.sglang_diffusion_utils.monkey_patches import ( patch_ltx2_disable_av_cross, + patch_ltx2_norm_out_fp32, patch_ltx2_rollout_cond_kwargs, patch_ltx2_trivial_attention_mask, ) @@ -66,6 +67,7 @@ def apply_ltx2_rollout_patches() -> None: patch_ltx2_rollout_cond_kwargs.apply() patch_ltx2_disable_av_cross.apply() patch_ltx2_trivial_attention_mask.apply() + patch_ltx2_norm_out_fp32.apply() def apply_env_selected_rollout_patches() -> None: diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_norm_out_fp32.py b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_norm_out_fp32.py new file mode 100644 index 00000000..55c6bbc6 --- /dev/null +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/patch_ltx2_norm_out_fp32.py @@ -0,0 +1,37 @@ +"""Match training's rounding points on the LTX2 output tail. + +Under autocast the training tail runs LayerNorm and the scale/shift modulation +in fp32 and rounds once at the proj_out matmul; rollout rounds to bf16 after +every op (~2e-3 on proj_out). Keep norm_out fp32 so the modulation promotes, +and cast at proj_out input like autocast does. +""" + +from __future__ import annotations + +import torch.nn.functional as F + + +def apply() -> None: + from sglang.multimodal_gen.runtime.models.dits import ltx_2 + + def fp32_norm_forward(self, x): + return F.layer_norm(x.float(), self.normalized_shape, self.weight, self.bias, self.eps) + + orig_init = ltx_2.LTX2VideoTransformer3DModel.__init__ + + def __init__(self, *args, **kwargs): + orig_init(self, *args, **kwargs) + for norm_name, proj_name in (("norm_out", "proj_out"), ("audio_norm_out", "audio_proj_out")): + norm = getattr(self, norm_name, None) + proj = getattr(self, proj_name, None) + if norm is None or proj is None: + continue + norm.forward = fp32_norm_forward.__get__(norm) + orig_proj_forward = proj.forward + + def proj_forward(x, _orig=orig_proj_forward, _proj=proj): + return _orig(x.to(next(_proj.parameters()).dtype)) + + proj.forward = proj_forward + + ltx_2.LTX2VideoTransformer3DModel.__init__ = __init__ From 255003302c2c5ee48a8b1dda4c75af1f38b64be9 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 4 Aug 2026 14:29:33 -0700 Subject: [PATCH 50/50] test(e2e): re-record the LTX-2.3 standard after the rollout parity patches --- .../test_ltx23_pickscore_grpo_4xGPU.json | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json b/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json index c97230ee..e648dbcf 100644 --- a/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json +++ b/tests/ci/fixtures/e2e_standards/test_ltx23_pickscore_grpo_4xGPU.json @@ -1,27 +1,27 @@ { "meta": { - "commit": "64d14300b551d673f858b2da7834a61aab48ff0e", + "commit": "b940e5ac5852f54d3e6ca7d709b8fe0e2cffef37", "source": "test_ltx23_pickscore_grpo_4xGPU.py" }, "metrics": { "rollout/reward/raw_mean": [ [ 0, - 0.6850860714912415 + 0.6851251125335693 ], [ 1, - 0.6817942261695862 + 0.6815934181213379 ] ], "rollout/reward/raw_median": [ [ 0, - 0.6864451766014099 + 0.6851017475128174 ], [ 1, - 0.6878968477249146 + 0.6879990100860596 ] ], "rollout/reward/raw_num_samples": [ @@ -37,65 +37,65 @@ "rollout/reward/raw_std": [ [ 0, - 0.04088980332016945 + 0.040446020662784576 ], [ 1, - 0.041681572794914246 + 0.042133599519729614 ] ], "train/grad_norm": [ [ 1.0, - 2.8766138711944222e-05 + 1.227360189659521e-05 ], [ 2.0, - 1.2603897630469874e-05 + 1.2525416423159186e-05 ], [ 3.0, - 1.8178370737587102e-05 + 2.1042627849965356e-05 ], [ 4.0, - 2.2844391423859634e-05 + 2.044663233391475e-05 ] ], "train/log_prob_mean_abs_diff": [ [ 1.0, - 3.625949223836263e-07 + 3.3155083656311035e-07 ], [ 2.0, - 7.146348555882772e-07 + 6.531675656636556e-07 ], [ 3.0, - 1.1622905731201172e-06 + 1.3715277115503948e-06 ], [ 4.0, - 1.274670163790385e-06 + 1.2914339701334636e-06 ] ], "train/log_prob_new_idx_0": [ [ 1.0, - -0.8474981039762497 + -0.8474980803827444 ], [ 2.0, - -0.8480802631626526 + -0.8480803271134695 ], [ 3.0, - -0.7732198803375164 + -0.7732197642326355 ], [ 4.0, - -0.7731634061783552 + -0.7731632639964422 ] ], "train/log_prob_old_idx_0": [