feat(ltx): consume rollout sigmas exactly and train with fp32 master weights

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

miles/backends/fsdp_utils/configs/ltx.py

@@ -13,18 +13,26 @@
 
 @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",)
     model_backend_path = "miles.backends.fsdp_utils.model_backend.MilesModelBackend"
     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.

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",

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:

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(

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: