From f142464088c4c3ea84ad0f0d3b45d7f1865ad452 Mon Sep 17 00:00:00 2001 From: rockdu Date: Fri, 7 Aug 2026 10:38:14 -0700 Subject: [PATCH 1/3] refactor(fsdp): family-declared input dtype policy at the model boundary --- miles/backends/fsdp_utils/actor.py | 47 ++++++------ miles/backends/fsdp_utils/configs/ltx.py | 9 +-- .../configs/train_pipeline_config.py | 2 + .../backends/fsdp_utils/input_dtype_policy.py | 50 ++++++++++++ .../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 -- .../fsdp_utils/test_input_dtype_policy.py | 76 +++++++++++++++++++ 8 files changed, 162 insertions(+), 54 deletions(-) create mode 100644 miles/backends/fsdp_utils/input_dtype_policy.py delete mode 100644 miles/backends/fsdp_utils/loss_hub/utils.py 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 d9430378..e8af2d5b 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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) diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index a4f77e7a..f4a2604c 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/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. diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index c392ea57..ecae05b4 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/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. diff --git a/miles/backends/fsdp_utils/input_dtype_policy.py b/miles/backends/fsdp_utils/input_dtype_policy.py new file mode 100644 index 00000000..a6790f3c --- /dev/null +++ b/miles/backends/fsdp_utils/input_dtype_policy.py @@ -0,0 +1,50 @@ +"""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.""" + 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 diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py index 00f787b6..876d69dc 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 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, diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 183bd47f..97932078 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 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: 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/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..f5f8f076 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_input_dtype_policy.py @@ -0,0 +1,76 @@ +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.input_dtype_policy 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_is_passthrough(): + 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.float32 + assert out_timesteps.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() + 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_cast_policy_casts_floats_only(): + latents, timesteps, pos_cond = _inputs() + out_latents, _, (out_pos, _, _) = apply_input_dtype_policy( + {"latents": "default", "cond": "default", "timestep": "fp32"}, + latents=latents, + timesteps=timesteps, + 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, 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 2f0a4eef4556c6fe4285639f964c636c67775058 Mon Sep 17 00:00:00 2001 From: rockdu Date: Fri, 7 Aug 2026 10:38:14 -0700 Subject: [PATCH 2/3] test(e2e): re-record the SD3.5 OCR standard for autocast boundary 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 ec23b50082d5f84ad82379a004bdf56271a732f4 Mon Sep 17 00:00:00 2001 From: rockdu Date: Fri, 7 Aug 2026 17:32:11 -0700 Subject: [PATCH 3/3] style(fsdp): clarify the input policy branches --- .../backends/fsdp_utils/input_dtype_policy.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/miles/backends/fsdp_utils/input_dtype_policy.py b/miles/backends/fsdp_utils/input_dtype_policy.py index a6790f3c..966d124f 100644 --- a/miles/backends/fsdp_utils/input_dtype_policy.py +++ b/miles/backends/fsdp_utils/input_dtype_policy.py @@ -26,25 +26,29 @@ def apply_input_dtype_policy( ) -> 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 _axis(key: str) -> torch.dtype | None: - axis = policy.get(key) - if axis is None: + 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 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] + 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 + return value # passthrough inputs, int masks, and non-tensors stay untouched 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 + 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), ) - return _cast(latents, latents_dtype), _cast(timesteps, timestep_dtype), out_conds