diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index d9430378..36386b64 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -39,6 +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, compile_precision, log_precision_summary, resolve_dtype from .sequence_parallel.plan import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -94,8 +95,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 @@ -128,6 +129,15 @@ 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 spec on clean FQNs (pre-LoRA, pre-FSDP). + compiled_precision = compile_precision( + model, + self.train_pipeline_config.precision_spec, + default_dtype=self._forward_dtype, + ) + if rank == 0: + log_precision_summary(component, compiled_precision, default_dtype=self._forward_dtype) + if args.use_lora: model = apply_lora(model, args, self.train_pipeline_config) @@ -143,6 +153,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), + compiled_precision=compiled_precision, ) checkpoint.broadcast_full_state_to_fsdp( model, @@ -512,22 +523,31 @@ 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) + # Boundary dtypes are family policy; op interiors stay autocast-managed. + 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, + sigmas=prepared.sigmas, + conds=(prepared.pos_cond, prepared.neg_cond, prepared.joint_cond), + default_dtype=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, - ) + with torch.autocast("cuda", dtype=forward_dtype, enabled=forward_dtype != torch.float32): + return train_pipeline_config.compute_noise_pred( + 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, + 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 +600,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,7 +625,14 @@ 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, + compiled_precision=None, +): from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard offload_policy = CPUOffloadPolicy() if cpu_offload else None @@ -619,33 +642,28 @@ 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}, precision wrap units={len(compiled_precision.wrap_units)}" ) - fsdp_kwargs = { - "mp_policy": MixedPrecisionPolicy( - param_dtype=param_dtype, - reduce_dtype=reduce_dtype, - ), - "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) - - for module in modules: - fully_shard(module, **fsdp_kwargs) - - fully_shard(model, **fsdp_kwargs) + 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( + param_dtype=policy_param_dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=False, + ), + "offload_policy": offload_policy, + "mesh": mesh, + } + + 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)) return model diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index a4f77e7a..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",) @@ -25,6 +31,8 @@ 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 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: self._height = args.diffusion_height @@ -91,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, @@ -117,41 +126,37 @@ 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 from ltx_core.utils import to_denoised - device = latents_input.device 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, ) - # 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/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 c392ea57..980a50e8 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]] = {} @@ -75,13 +76,16 @@ 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, ...] = () supports_cfg_training: bool = True # Rollout parity patch group applied by the engine (see monkey_patches; None = none). rollout_patch_group: str | None = None + # 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} # 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. @@ -102,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, @@ -110,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 00f787b6..3538d676 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 @@ -15,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], @@ -63,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 = ( @@ -74,28 +79,20 @@ def prepare_flow_grpo_batch( if use_cfg else None ) + # 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: - 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, 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 183bd47f..8243a36e 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,22 +37,16 @@ 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 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 - 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, @@ -128,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/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..4bd20fcf --- /dev/null +++ b/miles/backends/fsdp_utils/precision.py @@ -0,0 +1,226 @@ +"""Fine-grained weight-precision control for FSDP2, at module granularity. + +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. + +``compile_precision`` lowers the rules onto what FSDP2 can express: + + PrecisionSpec rules + | + (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 + +import logging +from dataclasses import dataclass +from fnmatch import fnmatch + +import torch + +logger = logging.getLogger(__name__) + +_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16} + + +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) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ModuleSel: + """Module selector; fqn and cls are globs over the module FQN and class name.""" + + 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: + """gather is a dtype name ("fp32"/"bf16"/"fp16") or "default", the run's default dtype.""" + + select: ModuleSel + gather: str + + +@dataclass(frozen=True) +class PrecisionSpec: + rules: tuple[Rule, ...] = () + + +# --------------------------------------------------------------------------- +# Compiler: spec -> FSDP2 lowering (per-module wrap units) +# --------------------------------------------------------------------------- + + +@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 + + +@dataclass +class CompiledPrecision: + 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 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] + # 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])] + + +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 compile_precision( + model: torch.nn.Module, + spec: PrecisionSpec, + *, + default_dtype: torch.dtype, +) -> CompiledPrecision: + """Resolve the spec against the (pre-LoRA, pre-FSDP) model into FSDP2 wrap units. + + 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. + 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} + hits = [0] * len(spec.rules) + + for mod_fqn, module in model.named_modules(): + parent_gather = gather_dtypes[_parent_fqn(mod_fqn)] + 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) + + 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") + 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): + if not hit: + raise ValueError(f"precision rule matched no module: {rule}") + 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.wrap_units)} extra wrap units" + ) + for unit in compiled.wrap_units: + logger.info(f"precision[{component}]: wrap {unit.fqn} @ {unit.param_dtype}") + + +# --------------------------------------------------------------------------- +# 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, + sigmas: torch.Tensor, + conds: tuple, + default_dtype: torch.dtype, +) -> 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. 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}") + + def _axis(key: str) -> torch.dtype | None: + axis = policy.get(key) + 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 _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(): + 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), _cast(sigmas, timestep_dtype), out_conds diff --git a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py index 2755e204..170ea7b4 100644 --- a/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py +++ b/miles/backends/sglang_diffusion_utils/monkey_patches/__init__.py @@ -59,11 +59,15 @@ 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, ) 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__ 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 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index e0c9c4ac..14d6df11 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. 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( "--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( @@ -1396,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 @@ -1492,6 +1505,12 @@ def set_default_diffusion_args(args) -> None: else: args.ref_mode = "none" + # --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: + args.sglang_dit_precision = args.precision_default_dtype + def miles_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args) @@ -1519,6 +1538,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-ltx23-sglang.sh b/scripts/run-diffusion-grpo-ltx23-sglang.sh index f6693627..ecbca076 100644 --- a/scripts/run-diffusion-grpo-ltx23-sglang.sh +++ b/scripts/run-diffusion-grpo-ltx23-sglang.sh @@ -80,9 +80,8 @@ 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 \ - --sglang-dit-precision bf16 \ --advantage-estimator grpo \ --globalize-reward-std \ --rm-type pickscore \ 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 \ 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..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": "5a1b7986af06e3d80368935c88acdd3dc6d98a40", + "commit": "b940e5ac5852f54d3e6ca7d709b8fe0e2cffef37", "source": "test_ltx23_pickscore_grpo_4xGPU.py" }, "metrics": { "rollout/reward/raw_mean": [ [ 0, - 0.6850860714912415 + 0.6851251125335693 ], [ 1, - 0.6819055080413818 + 0.6815934181213379 ] ], "rollout/reward/raw_median": [ [ 0, - 0.6864451766014099 + 0.6851017475128174 ], [ 1, - 0.6862333416938782 + 0.6879990100860596 ] ], "rollout/reward/raw_num_samples": [ @@ -37,65 +37,65 @@ "rollout/reward/raw_std": [ [ 0, - 0.04088980332016945 + 0.040446020662784576 ], [ 1, - 0.042159032076597214 + 0.042133599519729614 ] ], "train/grad_norm": [ [ 1.0, - 1.147376860899385e-05 + 1.227360189659521e-05 ], [ 2.0, - 1.245457497134339e-05 + 1.2525416423159186e-05 ], [ 3.0, - 2.025627691182308e-05 + 2.1042627849965356e-05 ], [ 4.0, - 1.9314091332489625e-05 + 2.044663233391475e-05 ] ], "train/log_prob_mean_abs_diff": [ [ 1.0, - 1.0107954343159993e-06 + 3.3155083656311035e-07 ], [ 2.0, - 1.188988486925761e-06 + 6.531675656636556e-07 ], [ 3.0, - 1.8222878376642864e-06 + 1.3715277115503948e-06 ], [ 4.0, - 1.761441429456075e-06 + 1.2914339701334636e-06 ] ], "train/log_prob_new_idx_0": [ [ 1.0, - -0.8474982244273027 + -0.8474980803827444 ], [ 2.0, - -0.8480804494271675 + -0.8480803271134695 ], [ 3.0, - -0.7732201963663101 + -0.7732197642326355 ], [ 4.0, - -0.7731628809124231 + -0.7731632639964422 ] ], "train/log_prob_old_idx_0": [ 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": [ 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..c8111c5c --- /dev/null +++ b/tests/fast/backends/fsdp_utils/_precision_wrap_worker.py @@ -0,0 +1,142 @@ +"""Gloo worker asserting the compiled precision plan really wraps under FSDP2 (2 ranks). + +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) + Rule(fqn="blocks.0.attn.norm_q", gather=default) + +so the tree and the dtype each module's params carry in the forward come out as: + + Net gather + ├── stem Leaf bf16 (no rule, wrapped by the root unit) + └── blocks + ├── 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 (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 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 + +DEFAULT_DTYPE = torch.bfloat16 + + +class Leaf(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 Norm(Leaf): + pass + + +class Attn(nn.Module): + def __init__(self): + super().__init__() + self.norm_q = Norm() + self.proj = Leaf() + + 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 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 + + +SPEC = PrecisionSpec( + rules=( + Rule(ModuleSel(cls="Norm"), gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.attn"), gather="fp16"), + Rule(ModuleSel(fqn="blocks.0.attn.norm_q"), gather="default"), + ) +) +EXPECTED_GATHER = { + "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") + world_size = dist.get_world_size() + mesh = init_device_mesh("cpu", (world_size,), mesh_dim_names=("dp_shard",)) + model = Net().to(torch.float32) # fp32 master + + compiled = compile_precision(model, SPEC, default_dtype=DEFAULT_DTYPE) + + 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 compiled.wrap_plan(model, list(model.blocks)): + fully_shard(unit.module, **fsdp_kwargs(unit.param_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/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 new file mode 100644 index 00000000..f47cb30a --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py @@ -0,0 +1,85 @@ +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) + 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, sigmas, pos_cond + + +def test_default_policy_is_passthrough(): + """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, sigmas, pos_cond = _inputs() + policy = {**DEFAULT_POLICY, "timestep": "default"} + _, 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(): + """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, + ) + 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(): + 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, + ) 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..3ff75bcf --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_precision_plan.py @@ -0,0 +1,256 @@ +"""Compiling PrecisionSpec rules into FSDP2 wrap units. + +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 + └── 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 + +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, PrecisionSpec, Rule, compile_precision + + +class Rope(nn.Module): + def __init__(self): + super().__init__() + 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.linear = nn.Linear(8, 8) + self.norm = nn.LayerNorm(8) + self.attn = Attn() + self.rope = Rope() + + +class Tiny(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([Block(), Block()]) + + +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) 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")} + + +def test_empty_spec_compiles_to_nothing(): + """No rules, so every node keeps the default and nothing is emitted. + + blocks.0 bf16 + ├── linear bf16 + ├── norm bf16 + ├── attn bf16 + │ └── norm_q bf16 + └── rope bf16 + """ + 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*", gather=fp32) + + blocks.0 bf16 + ├── linear bf16 + ├── norm [U] fp32 + ├── attn bf16 + │ └── norm_q [U] fp32 + └── rope bf16 + """ + 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(): + """Selecting by class name reaches the same two norms without naming any path. + + Rule(cls="*LayerNorm", gather=fp32) + + blocks.0 bf16 + ├── norm [U] fp32 + └── attn + └── norm_q [U] fp32 + """ + spec = PrecisionSpec(rules=(Rule(ModuleSel(cls="*LayerNorm"), gather="fp32"),)) + compiled = compile_precision(Tiny(), spec, default_dtype=torch.bfloat16) + assert set(_units(compiled)) == NORM_FQNS + + +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", gather=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 = Tiny() + spec = PrecisionSpec(rules=(Rule(ModuleSel(fqn="blocks.1"), gather="fp32"),)) + compiled = compile_precision(model, spec, default_dtype=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 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 blocks.1 + ├── norm [U] fp16 ├── norm [U] fp32 + └── attn └── attn + └── norm_q [U] fp16 └── norm_q [U] fp32 + """ + spec = PrecisionSpec( + rules=( + Rule(ModuleSel(cls="LayerNorm"), gather="fp32"), + Rule(ModuleSel(fqn="blocks.0.*norm*"), gather="fp16"), + ) + ) + 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, + "blocks.1.attn.norm_q": torch.float32, + } + + +def test_empty_module_sel_rejected(): + """A selector with neither fqn nor cls would silently match every module.""" + with pytest.raises(ValueError, match="needs fqn or cls"): + ModuleSel() + + +def test_every_node_of_a_nested_chain_wraps_bottom_up(): + """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 wrap order 3 + ├── linear fp16 (inherits blocks.0) + ├── norm fp16 (inherits blocks.0) + ├── 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 = Tiny() + 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"), + ) + ) + compiled = compile_precision(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, + } + assert _plan(model, compiled) == [ + ("blocks.0.attn.norm_q", torch.bfloat16), + ("blocks.0.attn", torch.float32), + ("blocks.0", torch.float16), + ("blocks.1", torch.bfloat16), + ] + + +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) + + 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 = 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 _plan(model, compiled) == [ + ("blocks.0", torch.float32), + ("blocks.1", torch.float32), + ("blocks", torch.float32), + ] + + +def test_paramless_module_gets_no_unit(): + """Buffers are never gathered, so pinning a buffer-only module lowers to nothing. + + Rule(cls="Rope", gather=fp32) + + blocks.0 + └── rope.freqs buffer -> no unit + """ + 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"), 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), + ] 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"]))