refactor(fsdp)(4/5): per-parameter weight-precision control for FSDP2 - #91
Draft
Rockdu wants to merge 8 commits into
Draft
refactor(fsdp)(4/5): per-parameter weight-precision control for FSDP2#91Rockdu wants to merge 8 commits into
Rockdu wants to merge 8 commits into
Conversation
…instead of ignored_params
…n_args and slim compiler output types
…on module docstring
Rockdu
changed the title
Rockdu
added a commit
to Rockdu/miles_diffusion
that referenced
this pull request
Aug 3, 2026
PR radixark#91 removed trainer-side forward input casts, so the LTX train forward began receiving an fp32 trajectory latent. dtype derivation in forward_velocity then sent fp32 positions into RoPE, while the sglang-d rollout computes RoPE from bf16-quantized coordinates: the pe tables diverge at rel ~0.5 in high-frequency bands (dump-verified on a paired LTX-2.3 run) and every attention output degrades from block 0. Casting the latent back to the autocast dtype at the model boundary restores the rollout-parity chain: bit-exact latent (bf16(x_fp32) equals the tensor rollout consumed), bf16 positions/context, and the original e2e reconstruction rounding. Verified on the same run that the fp32-sigma change works as intended: AdaLN timestep input now pairs fp32/fp32 at rel 6.2e-8 (was 1.8e-3), and blocks.0.attn1 input improves 17x (1.26e-2 -> 7.3e-4). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
A configurable precision-control system for FSDP2 training. Each model family declares per-tensor dtype intent as
PrecisionSpecrules in itsTrainPipelineConfig(in-repo Python, always active); a compiler (precision.py) lowers the per-tensor plan onto what FSDP2 can express. See the flow diagram in theprecision.pymodule docstring.Two dtype axes per rule, last match wins per axis:
master: resident dtype of the param/buffer (optimizer precision), lowered to per-tensor load-time castsgather: all-gather / forward dtype, lowered to nestedfully_shardgroups (one per gather dtype,reshard_after_forward=False) so overrides stay fully inside FSDP (DTensor params, FSDP grad reduction, DCP/offload unchanged)Compute dtype is deliberately not strictly managed: the trainer wraps the DiT forward in
torch.autocast(default dtype); op-level train/rollout gaps are left to the monkey-patch registry.Changes
miles/backends/fsdp_utils/precision.py: spec, compiler, lowering validation (whole-module coverage, single dtype per module, zero-hit rules rejected)--precision-default-dtype: one knob filling every dtype left unset (training forward/gather +--sglang-dit-precision); defaults filled inset_default_diffusion_args, conflicts rejected inmiles_validate_argscast_forward_inputs=False,cast_cond_to_dtypedeleted); DiT forward runs under autocast; LTX's inner autocast folded into the shared oneapply_fsdp2removed (autocast keeps recompute dtypes consistent)Tests
10 CPU unit tests for the compiler in
tests/fast/backends/fsdp_utils/test_precision_plan.py, registered tostage-a-cpu.Pending GPU validation (why draft)