Skip to content

fsdp: materialize clip_grad_norm's DTensor before logging - #35

Merged
zhihengy merged 1 commit into
mainfrom
fix/grad-norm-logging
Jul 13, 2026
Merged

fsdp: materialize clip_grad_norm's DTensor before logging#35
zhihengy merged 1 commit into
mainfrom
fix/grad-norm-logging

Conversation

@zhihengy

@zhihengy zhihengy commented Jul 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

clip_grad_norm_ on DTensor parameters returns a lazily-reduced partial norm; the logging path (.detach() → .mean().item()) leaked the local shard's value, under-reporting the train/grad_norm metric by ×√(n_shards) on every FSDP config (historical dp4 wandb curves are ×2 too low). Clipping itself was always correct — the clip coefficient is computed inside DTensor arithmetic — only the logged metric was wrong.

Fix: .full_tensor() before logging.

Found while cross-validating the two SP parameter placements in #21 (different shard counts → reported norms differed by √2 while instrumented per-param grads were bitwise identical). #21's validation claims depend on this landing first.

Expect wandb grad_norm curves to step up (×2 on the dp4 recipe) after this merges — historical values were under-reported, not the new ones inflated.

@zhihengy @claude
diffusion: materialize clip_grad_norm's DTensor before logging
41d9bbb 
clip_grad_norm_ returns a lazily-reduced partial-norm DTensor; logging
it without full_tensor() leaks the local shard's norm, under-reporting
grad_norm by sqrt(n_shards) (2x on the dp4 recipe). Clipping itself was
always correct — the coefficient is computed in DTensor arithmetic.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@zhihengy
zhihengy force-pushed the fix/grad-norm-logging branch from d2361fd to 41d9bbb Compare July 13, 2026 21:20
Rockdu
Rockdu approved these changes Jul 13, 2026
View reviewed changes

@Rockdu Rockdu left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

zhihengy added a commit that referenced this pull request Jul 13, 2026
@zhihengy @claude
Revert the clip_grad_norm logging materialization (split to #35)
94bc339 
Partial revert of 0420f3b: that commit also carried the dp_mesh -> fsdp_mesh
line (folded in during a rebase), which belongs to the SP placement and
stays. Only the DTensor import and the full_tensor materialization leave;
they land independently as #35, which this PR's cross-placement grad_norm
comparisons depend on.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@zhihengy zhihengy mentioned this pull request Jul 13, 2026
Merged
@zhihengy
zhihengy merged commit e833151 into main Jul 13, 2026
10 checks passed
zhihengy added a commit that referenced this pull request Jul 20, 2026
@zhihengy @claude
Revert the clip_grad_norm logging materialization (split to #35)
bf5448d 
Partial revert of 0420f3b: that commit also carried the dp_mesh -> fsdp_mesh
line (folded in during a rebase), which belongs to the SP placement and
stays. Only the DTensor import and the full_tensor materialization leave;
they land independently as #35, which this PR's cross-placement grad_norm
comparisons depend on.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
zhihengy added a commit that referenced this pull request Jul 21, 2026
@zhihengy @claude
Revert the clip_grad_norm logging materialization (split to #35)
36dd21b 
Partial revert of 0420f3b: that commit also carried the dp_mesh -> fsdp_mesh
line (folded in during a rebase), which belongs to the SP placement and
stays. Only the DTensor import and the full_tensor materialization leave;
they land independently as #35, which this PR's cross-placement grad_norm
comparisons depend on.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
zhihengy added a commit that referenced this pull request Jul 24, 2026
@zhihengy @claude
Revert the clip_grad_norm logging materialization (split to #35)
f51d293 
Partial revert of 0420f3b: that commit also carried the dp_mesh -> fsdp_mesh
line (folded in during a rebase), which belongs to the SP placement and
stays. Only the DTensor import and the full_tensor materialization leave;
they land independently as #35, which this PR's cross-placement grad_norm
comparisons depend on.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Rockdu added a commit that referenced this pull request Jul 25, 2026
@zhihengy @claude
@Rockdu @cursoragent
Feat: USP sequence parallelism (Ulysses x Ring) via SequenceParallelP…
90b6451 
…lan; Wan enabled (#21)

* diffusion: sequence-parallel config and Option B parallel state

FSDP shards on the dp mesh dim only; SP (ulysses x ring) gets its own
process groups plus sglang's _SP coordinator so USPAttention resolves
them. context_parallel_size stays as a backward-compatible alias.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: USP sequence parallelism for Wan training

Self-attention runs sglang's USPAttention; the sequence is sharded before
the first block and gathered after proj_out, so model outputs stay
full-sequence and loss/log_prob code is unchanged. Partial grads are
summed across the sp group after scaler.unscale_. Applied per component,
covering dual-DiT. Recipe gains SP_SIZE/ULYSSES_DEGREE/RING_DEGREE.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: align engine grouping span with rollout, assert full coverage

connect_rollout_engines grouped by raw rollout_num_gpus_per_engine while
rollout.init_rollout_engines uses min(per_engine, num_gpus_per_node);
a mismatch orphans trailing ranks that fail later with an opaque error.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: SP test suite (mesh, init smoke, parity, weight sync, import guard)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: SP guardrail runner + fp32 mode for grad-sync parity

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: guardrail runner cleanup preamble + expert-selection override

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: drive SP shard/gather hooks from the model's diffusers _cp_plan

The shard/gather placement (rope outputs, first-block input, proj_out
output) was hard-coded to Wan's module topology; diffusers 0.37 declares
the same contract per model as _cp_plan, so consume that instead. The
attention processor remains Wan-specific and is now guarded explicitly.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: guardrail runner supports rollout-data save/replay for same-data band comparison

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: own the differentiable USP operators, drop the sglang training dependency

RL needs a differentiable attention in the trainer, not in the rollout
engine; importing sglang's USPAttention leaked that requirement into a
fork the engine had to carry. sp_ops.py holds the Ulysses all-to-all,
ring attention (torch ring templates), and SDPA local attention with the
same layout, so training numerics are unchanged and sglang needs no
patches. The trainer's only remaining sglang import is the weight-sync
checksum.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: first-principles SP cleanup

- drop context_parallel_size and its CP->SP aliasing: the old assert==1
  meant no config with cp>1 could ever have existed
- drop the num_heads param from validate_sp_config and the always-None
  getattr feeding it: the real guard lives in apply_sequence_parallel
- degree-1 ulysses/ring dimensions get no process groups (None = local),
  and the unused per-rank singleton tp_group is gone
- move the sequence shard/gather collectives into sp_ops so all
  differentiable SP collectives live in one module
- remove the dead --fwd-only flag (ring backward works)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: fsdp_shard_mode dp_sp — shard parameters over the dp x sp mesh

Aligns with torchtitan/VeOmni/DeepSpeed-Ulysses: FSDP shards over the
flattened dp x sp mesh (default), halving/quartering per-rank parameter
and master-state memory under SP (measured 54.4 -> 27.2 GB/rank at
dp2xsp2). Gradient semantics live in the sequence gather's backward
(sum over the sp group), so the explicit cross-sp grad all-reduce is
skipped; fsdp_shard_mode=dp keeps the replicated placement as a
validation anchor. Guardrail preamble no longer kills by process name;
it refuses to start when this run's GPUs are occupied.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: materialize clip_grad_norm's DTensor before logging

clip_grad_norm_ returns a lazily-reduced partial-norm DTensor; logging
it without full_tensor() leaks the local shard's norm, under-reporting
grad_norm by sqrt(n_shards) (2x on the dp4 recipe). Clipping itself was
always correct — the coefficient is computed in DTensor arithmetic.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: SequenceParallelPlan — per-family SP declaration behind ModelBackend

A model family opts into SP by declaring a SequenceParallelPlan (boundaries,
attention installer, head count) instead of the engine hard-coding Wan:

- ModelBackend.sequence_parallel_plan assembles the plan; the diffusers
  default reads the model's _cp_plan and the family config's
  apply_sp_attention hook. Native backends override the method wholesale.
- WanUSPAttnProcessor moves to configs/wan2_2.py; the engine layer
  (sp_attention/sp_ops/sp_mesh/parallel) no longer references any model class.
- The isinstance whitelist becomes capability checks: no _cp_plan or no
  apply_sp_attention raises with the missing declaration named.

Validated: CPU suites (88 passed) + full SP GPU suite (21 bands, 66 checks).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* style: pre-commit formatting across sp files

First full pre-commit pass on this branch (the job never ran while the PR
was based on feat/wan): isort import merges, black line wraps, ruff
Callable-import modernization, unused-import removal. No code changes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: fail-fast SP validation — driver-side support check, wildcard and key-resolution guards

Adversarial review of the plan seam surfaced four holes no numerics test
covers; all four now fail loudly at the earliest point the facts exist:

- validate_sp_support (driver arg validation, zero weight loads): rejects
  sp>1 with an explicit --fsdp-attention-backend (the SP installer replaces
  every processor, so backend selection and the deterministic-flash patch
  can never take effect), a model backend without sequence_parallel_plan,
  or a diffusers family without apply_sp_attention — previously these
  surfaced only after full load + FSDP sharding of every component.
- plan construction rejects wildcard _cp_plan boundaries (e.g. QwenImage's
  transformer_blocks.*) the hook installer cannot expand yet, instead of a
  bare AttributeError deep in the getattr walk.
- boundary key resolution reads the forward signature from the
  get_base_model()-unwrapped module (PeftModel exposes (*args, **kwargs))
  and raises at install time for unmappable keys instead of silently
  skipping the split.

Validated: 95 CPU tests (7 new), full SP GPU suite (21 bands, 66 checks),
grad-sync --lora bands for both shard modes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* tests(sp): determinism smoke — SP attention path bitwise under deterministic mode

Validates that deterministic_mode's guarantees survive USP: forward+backward
twice on identical inputs must produce bitwise-equal outputs and grads for
both the ulysses (SDPA local) and ring (aten flash fwd/bwd) paths, with
use_deterministic_algorithms(True, warn_only=False) engaged — warn_only=False
also asserts no op on the path is registered nondeterministic, confirming
the aten deterministic gate covers the ring backward torch templates call.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: parameter placement is Option A unconditionally — Option B demoted to a test anchor

FSDP always shards over the flattened dp x sp mesh; SP is invisible to it
(same parameter account as sp=1). Placement is not a user decision, so the
config surface shrinks to --sequence-parallel-size + --ulysses-degree
(ring = sp / ulysses, derived):

- fsdp_shard_mode and ring_degree leave arguments; the mesh selection in
  parallel.py is branch-free (comment notes the HSDP caveat: flatten the
  shard axes, which today happen to be the whole world).
- _all_reduce_sp_grads leaves the actor; the sequence gather's sum_grad
  backward is the only gradient mechanism in production.
- The sp-replicated Option B placement survives verbatim in tests/sp as a
  validation anchor: grad-sync/weight-sync build it from the exported
  dp_mesh and apply_sequence_parallel(..., sum_grad=False), so the
  cross-placement parity bands still exercise production machinery.

Validated: 95 CPU tests, full SP GPU suite (66 checks), grad-sync --lora
for both placements, determinism smokes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: ring_degree is derived, not configured — drop it from every input surface

ring = sp / ulysses is fully determined, so taking it as input carried zero
information; its one historical value (cross-checking contradictory u,r
pairs) died with the config-surface removal. sp_mesh's pure functions now
take (sp, ulysses) only and return the derived ring; the u*r==sp product
check collapses to a plain sp % u divisibility check. Test CLIs express
ring topologies as (sp, ulysses): u2r2 = --sp 4 --ulysses 2.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* docs(sp_mesh): note the sp_subgroups <-> locate_rank inverse relation

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: default SP attention intercepts at diffusers' dispatch_attention_fn

A model family no longer writes an attention processor: the default
apply_sp_attention wraps the modeling module's dispatch_attention_fn (found
through the MRO — fully_shard swizzles the class into torch's fsdp module)
and reroutes self-attention call sites to sp_ops.usp_attention. Models flag
self- vs cross-attention themselves (upstream passes parallel_config only
when encoder_hidden_states is None), so the model's own processors run
untouched and WanUSPAttnProcessor (the 80-line copy) is deleted.

The wrapper fails loudly outside the validated envelope (mask / dropout /
is_causal / custom scale); models whose modules never imported
dispatch_attention_fn are rejected at install with instructions to provide
a family override. Kernel choice stays pinned inside usp_attention.

Validated: parity vs the stock diffusers processors is bitwise for forward
and input grads (fp32 band exact); full GPU suite 77 checks; 10-step real-RL
runs at sp=1 / u4 / u2r2 / u1r4 — per-step train-vs-rollout canary flat
across all topologies (no rounding accumulation), rewards and clipfrac in
the sp=1 band.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Revert "diffusion: align engine grouping span with rollout, assert full coverage"

This reverts commit bc44519.

* Revert the clip_grad_norm logging materialization (split to #35)

Partial revert of 0420f3b: that commit also carried the dp_mesh -> fsdp_mesh
line (folded in during a rebase), which belongs to the SP placement and
stays. Only the DTensor import and the full_tensor materialization leave;
they land independently as #35, which this PR's cross-placement grad_norm
comparisons depend on.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* tests: keep the SP suite local — it lands in a separate test-infra PR

The full validation suite (mesh invariants, attention/grad/weight-sync
parity, determinism smokes, RL guardrail runner) stays out of this PR by
maintainer decision; every band was run locally against this exact tree
(results recorded in the PR description). The files remain in the working
tree for local use and will be contributed with CI registration in a
follow-up PR.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* refactor: sp_plan.py owns the plan contract; sp_attention.py is dispatch-only

SequenceParallelPlan, the boundary-hook interpreter and apply_sequence_parallel
move to sp_plan.py; sp_attention.py keeps only the dispatch-level attention
installer. The two are siblings over sp_ops — the plan holds attention as an
opaque callable, so neither imports the other, and the file layout now matches
the dependency structure (model_backend imports the contract, the family-config
default imports the installer). Pure code motion, no behavior change.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: unify cp_*/sp_* in ParallelState — sp is the only sequence axis

cp_rank/cp_size/cp_group/dp_cp_group had no consumers outside the
sp-alias wiring; the dp x sp composites rename to dp_sp_* to match the
flattened FSDP mesh dim.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* diffusion: pin ring templates to their torch>=2.11 home, flag the experimental dependency

The CI image (sglang v0.5.12) pins torch==2.11.0, where torch moved the
private ring-attention templates to _context_parallel/_attention.py and
the back-compat stub re-exports only the forward one. Import the 2.11
home directly; older torch fails loudly at first ring use.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* style: isort on sp_ops.py

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* refactor(sp): make sequence parallel plan model-owned

* feat(sp): honor attention backend under pure Ulysses

* refactor(sp): trim test-only paths and harden validation

* refactor(sp): align validation with ownership boundaries

* refactor(sp): group sequence-parallel implementation

* test(sp): track attention parity and determinism

* sp: install USP via Wan attention processor; ring kernel follows the attention backend

* Revert "sp: install USP via Wan attention processor; ring kernel follows the attention backend"

This reverts commit c957a13.

* sp: ring attention honors --fsdp-attention-backend via RING_KERNELS

* refactor: signature bind

* refactor: absorb sp attention installation into ModelBackend

* refactor: absorb support sp as backend attribute

* style: format model backend

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: rockdu <kangrdu@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Rockdu Rockdu Rockdu approved these changes
@guapisolo guapisolo Awaiting requested review from guapisolo guapisolo is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@zhihengy @Rockdu