refactor(fsdp)(4.3/5): family-declared input dtype policy at the model boundary

miles/backends/fsdp_utils/actor.py

@@ -35,6 +35,7 @@
     DiffusionUpdateWeightFromTensorLoRAIPC,
 )
 from .ema import EmaShadow
+from .input_dtype_policy import apply_input_dtype_policy
 from .loss_hub import DiffusionLossContext, flow_grpo_loss_formula, prepare_flow_grpo_batch
 from .lr_scheduler import get_lr_scheduler
 from .metrics import new_metric_buffer
@@ -512,22 +513,29 @@ 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, (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:
-            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,
+                    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()
 
@@ -626,23 +634,16 @@ def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None, no_split_modules
     )
 
     fsdp_kwargs = {
+        # input_dtype_policy owns boundary casts; autocast owns compute and keeps grad-ckpt recompute consistent.
         "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)
-
     for module in modules:
         fully_shard(module, **fsdp_kwargs)
 

miles/backends/fsdp_utils/configs/ltx.py

@@ -25,6 +25,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
@@ -129,7 +131,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 +149,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.

miles/backends/fsdp_utils/configs/train_pipeline_config.py

@@ -82,6 +82,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
+    # Model-boundary input dtypes (see 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.

miles/backends/fsdp_utils/input_dtype_policy.py

@@ -0,0 +1,54 @@
+"""Family-declared dtypes for the model-boundary inputs (latents, cond, timestep).
+
+The trainer does not hard-cast its forward inputs; each family declares, per input, a dtype name
+("fp32"/"bf16"/"fp16"), "default" for the run's forward dtype, or None to pass the rollout dtype
+through. The boundary dtype is what element-wise ops see before any weight is involved, so it must
+match what the family's sglang-d pipeline feeds the DiT for log-prob alignment; compute inside the
+model is owned by the trainer's autocast (see actor.apply_fsdp2).
+"""
+
+from __future__ import annotations
+
+import torch
+
+_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}
+
+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 float boundary inputs per family policy ("default"/dtype name/None=passthrough);
+    autocast alone would leave element-wise ops running at the raw input dtype."""
+    # A typo'd key would silently mean passthrough.
+    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 _dtype(key: str) -> torch.dtype | None:
+        dtype_name = policy.get(key)
+        if dtype_name is None:  # passthrough: keep whatever dtype rollout handed us
+            return None
+        if dtype_name == "default":  # the run's forward dtype
+            return default_dtype
+        if dtype_name not in _DTYPES:
+            raise ValueError(f"input_dtype_policy[{key!r}] has unknown dtype {dtype_name!r}")
+        return _DTYPES[dtype_name]
+
+    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  # passthrough inputs, int masks, and non-tensors stay untouched
+        return value.to(dtype)
+
+    cond_dtype = _dtype("cond")
+    return (
+        _cast(latents, _dtype("latents")),
+        _cast(timesteps, _dtype("timestep")),
+        tuple(cond and {key: _cast(value, cond_dtype) for key, value in cond.items()} for cond in conds),
+    )

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

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 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:

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": [