Engine and model-layer refactor#70
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors PithTrain's state management by replacing context managers with process-global context modules, updates the 5-stage decoder layer protocol, and introduces support for variable-length sequence packing with token-weighted loss. It also renames the FSDP integration tests to test_dualpipev. The review comments correctly identify critical runtime bugs where torch.distributed.broadcast is called with the invalid argument group_src instead of src in both pithtrain/dualpipe/dualpipev.py and pithtrain/tasks/pretrain_lm.py.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
As part of the refactor, we've bumped cutlass-dsl >= 4.6.0 and get rid of the hack that avoids redundant |
8c05124 to
972e712
Compare
|
@claude review |
Code reviewReviewed with two independent CLAUDE.md/AGENTS.md-compliance passes and two independent bug-hunting passes, then adversarially validated every flagged issue. No logic bugs were found — the
Note: inline PR comments couldn't be posted in this run (the inline-comment tool wasn't available and |
989428c to
8e882e3
Compare
Fixed. |
Fixed. |
a09513c to
94600e2
Compare
|
@claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # Token-weighted reduction. The criterion returns per-micro-batch loss SUMS, so gradients | ||
| # accumulate d(sum loss); dividing by the total non-ignored token count yields the correct | ||
| # token-mean regardless of how tokens split across micro-batches. The count lives on pipeline | ||
| # rank 0 (the only rank with labels), so broadcast it to every stage for the gradient scale. | ||
| num_tokens = torch.ones((), device=device) | ||
| if distributed.pp_rank == 0: | ||
| num_tokens.fill_((global_labels != -100).sum().clamp_min(1)) | ||
| if distributed.pp_size > 1: | ||
| torch.distributed.broadcast(num_tokens, group_src=0, group=distributed.pp_group) |
There was a problem hiding this comment.
Latent correctness gap: num_tokens is a per-rank local count, not the "global" count the comment claims — biased once packed/variable-length training goes live.
num_tokens is (global_labels != -100).sum() on this (dp, cp, ep) coordinate's pp-rank-0, broadcast only along the PP group. It is never all-reduced across DP/CP/EP. The final token-mean gradient is therefore correct only because FSDP2 reduces gradients with ReduceOp.AVG over the FSDP mesh (dp,cp,ep for non-expert params, dp,cp for experts — modules/training.py:341-346) and every rank's token count is equal:
AVG_r(dS_r / N) == (Σ_r dS_r) / (R·N) == Σ dS_r / Σ N_r ✅ only when all N_r are equal.
For the live fixed-sequence-length pretraining path all N_r are equal, so this matches main exactly and the wandb parity curves are expected. But this PR's headline feature is the variable-length/packed interface, and once tokens-per-rank differ:
AVG_r(dS_r / N_r) ≠ Σ dS_r / Σ N_r
the gradient (and the logged loss, via the sum/cp_size on lines 390-391) becomes biased by the per-rank imbalance. The fix is to all-reduce the count over the same mesh FSDP reduces on before dividing, rather than relying on equal counts.
At minimum, the "global non-ignored token count" wording here and in criterion (lines 116-118) overclaims — it's local — and will mislead whoever wires up the packed data pipeline into thinking the scaling is already correct for unequal counts.
Correctness ReviewTraced the engine (execution/overlap/dualpipev), the 5-stage layer refactor + Verified sound (called out because these are the easy places to get a "forward-identical, gradients-wrong" refactor bug):
One correctness gap worth addressing (latent, not live) — see inline comment on Evidence assessment: |
| torch.distributed.broadcast(width, group_src=0, group=self.pp_group) | ||
| if int(width) == 0: |
There was a problem hiding this comment.
setup_cu_seqlens runs at the start of every DualPipeV.step(), including the live dense-pretraining path (cu_seqlens=None). When pp_size > 1 this now issues a per-step PP-group broadcast(width) followed by int(width), which forces a host↔device sync (.item()-equivalent) before the pipeline warmup can launch.
The cost is small — one scalar collective + one sync per step (not per micro-batch or per layer), so it's a fraction of a hundreds-of-ms step — but it's a genuinely new per-step sync + collective on the hot loop, added to support a feature (packed sequences) that no entrypoint drives yet. Because it precedes the whole pipeline schedule, the sync effectively acts as a PP-group barrier at each step boundary.
Worth a note rather than a blocker: consider skipping the width broadcast when the framework knows the run is dense (e.g. a config flag set once at setup), so the live path pays nothing until packed data is actually wired up.
Performance ReviewScope: judged whether this refactor holds training step-time and peak memory vs. No concrete throughput/memory regression found in the mechanics. The comm-overlap structure is preserved: dispatch/combine still run on 1. Evidence gap (primary). The validation is loss-parity wandb curves only. For a PR that restructures the overlap engine, the FP8 GEMM path, and the per-step reduction, loss parity does not establish that step-time and peak memory are unchanged — those are exactly the axes this framework exists to optimize, and they can regress with the loss curve still matching bit-for-bit. The step logger already emits 2. New per-step sync + collective on the live dense path (minor). pith-train/pithtrain/dualpipe/dualpipev.py Lines 185 to 186 in 94600e2 Neither point blocks the refactor. Item 1 is the one I would want addressed before merge given the surface area touched. |
Consistency ReviewThis refactor renames the decoder-layer stages ( The primary prose docs — AGENTS.md, README.md, 1.
2. 3. 4. 5.
|
Fixed! |
…ed with its intention
…w+sink Sibling-model cosmetic/cleanup pass (E), built to the caf8f50 end-state (last model-cleanup commit before variable-seqlen): - gpt-oss forward compute rewrite + torch.compile on stage1 (caf8f50, 22ef069); FA4 wrapper gains sliding-window (window_size) and attention-sink (learnable_sink) support that gpt-oss needs. - single-config pass-through construction across the four models (bf2140e). - unified modeling docstrings linking the HF reference models (32cd343); further modeling cleanup (ce9affb). Engine (execution/overlap/dualpipev/pretrain_lm) untouched. BF16 gate (deepseek-v2-lite, pp2/ep2, 32 steps): lb tracks base (max |Δ|=0.015, same 1.06->1.26->1.10 arc); CE tracks (6.436 vs 6.424). NOTE: deepseek changes only 2 lines here, so this gate validates the shared files (FA4/interface/contexts) and deepseek; the gpt-oss/qwen compute rewrites are not numerically exercised by a deepseek run. Variable-seqlen and the routing-info rename remain deferred.
Variable-length / packed-sequence support (8de11e2): cu_seqlens threaded through the pipeline as an explicit forward argument (sibling of rotary_posemb) into attention (flash_attn_varlen_func) and the per-document position reset; DualPipeV broadcasts it along the PP group. Loss becomes token-weighted: criterion returns the summed CE and train_step divides every gradient by the global non-ignored token count. Plus the MoERouting -> RoutingInfo rename (7f37007). NOTE: this batch diverged the lb-loss gate on its own (router gradient scaled ~num_tokens too small because the separately-injected load-balance gradient does not flow through criterion, yet is hit by the 1/num_tokens grad scale). Fixed in the immediately following commit; committed separately to record the regression and its fix distinctly.
The lb gradient is injected at the gate via MoELoadBalanceLossInjector, not routed through criterion, so the token-weighted-loss change (summed CE + 1/num_tokens gradient scaling) shrank it by ~num_tokens and the router stopped balancing (lb-loss climbs instead of decaying). Weight the injected lb_loss by the micro-batch token count (topk_weight.shape[0]) so the uniform 1/num_tokens scale normalizes it correctly -- equal to the previous 1/accumulate_steps mean for equal-length micro-batches. The returned/logged lb_loss is unchanged. Applied to all four gates (deepseek-v2, qwen3, qwen3.5, gpt-oss).
…ackward torch._dynamo.allow_in_graph(MoELoadBalanceLossInjector) at module scope. The gpt-oss router's forward is @torch.compile(fullgraph=True) and calls MoELoadBalanceLossInjector.apply inside it; without allow_in_graph Dynamo traces into the custom autograd Function and can drop its injected backward (the deepseek/qwen gates are not compiled, so this only bites gpt-oss). Also folds the dualpipev cu_seqlens line-wrap (640e471, cosmetic). Orthogonal to the load-balance gradient-scaling fix: that changes the injected value, this preserves the injected backward across compilation.
Pull the last deferred non-framework changes so 0709 equals the full 0630 codebase refactor plus the load-balance gradient-scaling fix, and nothing else: - pyproject.toml: flash-attn 4 b19 -> b20; drop the nvidia-cutlass-dsl-libs-base override in favor of prerelease="allow" + linux environment pin (78576f6). - deepseek-v2-lite/config.json: drop moe_layer_freq / norm_topk_prob / scoring_func (unused by the model) (5ac3a99). - test formatting (ruff) + agent-guide doc text (GroupLinear -> GroupedLinear, layers/ -> operators/, drop ConvertCheckpointCtx). After this, `git diff 0709 0630-codebase-refactor` is exactly the 4-gate load-balance fix.
deepseek_v2_lite.py -> deepseek_v2.py (theme E rename) and the layer stage methods forward_attn/forward_mlp/forward_aggregate -> forward_stage1/forward_stage3/forward_stage5. Both the supported-models line and the MLA note pointed at the old names.
Expand the paired semicolon-separated annotations in contexts/distributed.py to one per line and remove the now-unneeded E702 per-file-ignore. Reduce the three context modules to a single-line docstring; the access contract and usage examples are documented once in AGENTS.md instead.
Add a Runtime Contexts section (the pithtrain.contexts access contract: import the module, read fields off it, populated once by setup_*). Fix two stale spots from the contexts refactor: the training loop uses setup_logging/ setup_distributed/setup_training, not the old *_context managers, and tasks no longer carry a per-task <Task>Ctx object.
… a result, we drop the bf16 defaults
…ughout all decoder layers
…e-training and we shouldn't specify the max position embeddings from config since the pre-training will overwrite with sequence length ; then the model.json from benchmarks was moved ot the config.json inside pretrain examples for the qwen3.5 model
…file with the interface file.
… one-line md paragraphs
…rd_stage1/3/5, LayerRecord
… in capture/validate scripts
…oved layer_partition
529c3ec to
dc93ac5
Compare
2 participants
Engine and model-layer refactor
A cleanup of the engine and model layer. The training recipe is unchanged; the work removes plumbing that only forwarded arguments around and tightens the interfaces between the models, the pipeline, and the runtime.
Changes
Runtime contexts. Process-global state (distributed groups, the model and optimizers, the logger) lives in per-concern context modules, each filled in once at startup by a
setup_*function and read directly where needed. This replaces the per-taskCtxobjects that were threaded through every call and the*_contextmanagers that wrapped setup.Config passed down, not expanded early. Modules take the transformer config object and read the fields they need locally, instead of callers unpacking scalars and threading them through constructors. Fields are pulled out only where a value genuinely differs from the config (an override or a derived size), which keeps each module's dependencies visible at its own definition.
Decoder-layer stage naming. The layer's three compute entry points are renamed to
forward_stage1/3/5, matching the pipeline's five-stage attention → dispatch → MLP → combine → aggregate split (stages 2 and 4 are the engine-driven all-to-alls). The routing metadata that stage 1 produces and stage 5 consumes now travels as one named record instead of a bare tuple.Engine consolidation. The autograd wrappers and dispatch/combine helpers moved in alongside the stage implementations, so the forward/backward flow reads in one place. The overlap loop and layer partitioning were simplified in the same pass.
Variable-sequence-length interface. The pipeline can carry per-document boundaries through to attention and the position embedding, the groundwork for packed / variable-length (SFT) training. The residual stream stays batch-of-one BSHD; only attention sees the flattened view, and with no boundaries set the behavior is identical to before. No entrypoint produces packed batches yet (the corpus-packing data pipeline is a follow-up); the framework side is checked with synthetic boundaries in the multi-GPU test.
Cleanups. Each model defines its own stage methods and keeps its helpers local to its file, so it reads top to bottom on its own; DeepSeek dropped its "lite" naming for the base family; run logs unified under one directory; the multi-GPU test renamed; wandb off by default in examples; docs updated.
Validation
Compared against the pre-refactor baseline on two models spanning different model families, parallelism meshes, and load-balance strategies, both at a constant learning rate. The base-vs-feature load-balance loss curves are on wandb:
The curves match: load-balance loss rises, turns over, and decays on both, and cross-entropy agrees within bf16 run-to-run noise.