Repository navigation
[DeepSeek V3] Fuse general-shape sequence loss forward and backward - #5552
sanketpurandare wants to merge 1 commit into
Conversation
meta-cla
Bot
added
the
CLA Signed
87d562a to
9f34a27
Compare
sanketpurandare
changed the title
## Agent details (AI-generated) Adapted from Yidi Wu (@ydwu4), #5259, revision 25dafa9. Preserve original attribution and license; this change includes integration and generalization, not just a port. Keep provided-map/L1 and unbiased top-k objectives distinct. Preserve native coefficient/token normalization, gradient injection, metric accumulation and DP-local training ownership. Padding contributes no counts, score sums or gradients; unbiased selection is independent of biased/group-limited routing. General fused CUDA kernels support positive T/E and configured top-k, while the aligned production geometry keeps the tuned CuTe launch. Derive scratch and native CTA partitioning from shapes and device resources. Explicitly fold block.y partials in native order, independent of Triton layout; preserve unaligned score-base offsets and scalar clamp-min's strict boundary derivative. Temporary normalized scores/maps and partial buffers are not autograd saved state. Save only original scores, frequencies and an optional padding mask. The general implementation does not fall back to eager ATen for shape/tails. ## Test Plan (AI-generated) python -m pytest tests/unit_tests/cpu/test_fused_dsv3_router.py tests/unit_tests/gpu/test_fused_dsv3_router.py tests/unit_tests/cpu/test_fused_dsv3_seqwise_loss.py tests/unit_tests/gpu/test_fused_dsv3_seqwise_loss.py -q pre-commit run --all-files 86 combined CPU/GPU tests passed, covering exact gradients/losses, signed and tiny scores, clamp equality, all-padding rows, odd experts, native CTA-split boundaries, offset1/2/3 storage, FakeTensor metadata and fullgraph/capture/replay. The benchmark's post-hot seed14536 failure is preserved as an exact regression test; fixing native reduction order removed its one-ULP raw-loss mismatch. Strict relevant docs passed. All non-type lint hooks passed; only unchanged baseline type findings remained. Matched graph-replay F+B general cases beat the same native reference by 2.662x (T8193/E64/K6), 3.064x (T4097/E512/K8) and 2.517x (T33/E129/K8), with exact reference checks and separate eager GPU trace collection. The production path retains the same kernels; latest median20.933us versus 20.701us before overlaps the observed trial spread. A longer final matched hot-path comparison and complete PP2/PP4/memory gates remain pending. No distributed-throughput or full-model convergence claim follows here.
9f34a27 to
15a499b
Compare
pytorch-bot
Bot
added
ciflow/b200
|
@claude review |
|
Claude finished @tianyu-l's task in 4m 55s —— View job Verdict: changes requested (minor). The design is sound, but please remove the duplicated CuTe helpers, fix the tuned backward so it uses the same unbiased denominator as its forward, and widen GPU test gating so CI on B200 actually runs most of the new tests. See the 6 inline comments. Review detailsWhat's good
Inline findings
PR description
Not verifiedI couldn't run anything: running pytest/python here needs approval, and there are no GPUs. I did not check the bitwise-parity claims (ATen reduction-order emulation, clamp-min boundary derivative, top-k tie-breaking) beyond reading the code. I'm relying on the GPU test matrix described in the PR. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| @cute.jit | ||
| def add(a, b): | ||
| return arch.inline_ptx( | ||
| "add.rn.f32 {$w0}, {$r0}, {$r1};", | ||
| write_only_types=[Float32], | ||
| read_only_args=[a, b], | ||
| ) | ||
|
|
||
| @cute.jit | ||
| def mul(a, b): | ||
| return arch.inline_ptx( | ||
| "mul.rn.f32 {$w0}, {$r0}, {$r1};", | ||
| write_only_types=[Float32], | ||
| read_only_args=[a, b], | ||
| ) |
There was a problem hiding this comment.
Duplication: this block re-defines add, mul, div, choose, load4, store4, row_load, row_store, cluster_sync, remote_load, remote_store and launch. They're character-for-character copies of fused_dsv3_router.py:438-574, and the line above already imports select_eight from that module. Please import them as well (launch only differs by the missing bfloat16 entry in types, and keeping that entry is harmless). That's about 150 fewer lines here, and the two kernels can't drift apart on PTX helpers that both rely on for bitwise parity. Only fma, absf, clamp_norm, warp_sum, row_sum, row_reciprocal, fast_range and row_divide are new. If the router has the same helpers inline, those could move to the shared module too.
| magnitudes = cute.make_rmem_tensor(8, Float32) | ||
| for i in cutlass.range_constexpr(8): | ||
| magnitudes[i] = absf(items[i]) | ||
| norm = row_sum(magnitudes) | ||
| if cutlass.const_expr(unbiased): | ||
| denominator = add(norm, Float32(1e-20)) |
There was a problem hiding this comment.
Forward and backward use different unbiased denominators. In unbiased mode the tuned backward builds its denominator from row_sum(magnitudes) + 1e-20 (abs-sum). Three other places use the plain sum scores.sum(-1) + 1e-20:
- the tuned forward (
forward_kernel, L723-724:row_sum(items)) - the generic backward (
_generic_seqwise_backward, which passesscoreswhenUNBIASED) - the native
UnbiasedTopKLoadBalanceLoss
Sigmoid scores are nonnegative, so this doesn't show up today. But FusedDSv3SeqwiseLossFunction accepts signed input. Your generic tests even include a signed + unbiased=True case, and nothing tests that combination on the tuned path. With signed scores at 4096x256 on GB300, this kernel's gradient would not match its own forward. Suggest matching the forward:
norm = row_sum(items) if cutlass.const_expr(unbiased) else row_sum(magnitudes)(then magnitudes is only needed when the mode is biased).
| class FusedDSv3SeqwiseLoss(MicrobatchWiseLoadBalanceLoss): | ||
| """Specialize the loss while retaining native scaling and metric state.""" | ||
|
|
||
| @dataclass(kw_only=True, slots=True) | ||
| class Config(MicrobatchWiseLoadBalanceLoss.Config): | ||
| top_k: int | None = None | ||
| """None uses the provided map; a positive value selects unbiased top-k.""" | ||
| require_fusion: bool = False | ||
|
|
||
| def __post_init__(self): | ||
| self.sharding_config = _with_local_counter_sharding(self.sharding_config) | ||
| if self.top_k is not None and self.top_k <= 0: | ||
| raise ValueError("top_k must be positive") | ||
|
|
||
| def __init__(self, config: Config): | ||
| self.top_k = config.top_k | ||
| self.require_fusion = config.require_fusion | ||
| super().__init__(config) | ||
| self.register_buffer( | ||
| "arrival_counter", torch.zeros(1, dtype=torch.int32), persistent=False | ||
| ) | ||
|
|
||
| @property | ||
| def metric_name(self) -> str: | ||
| return ( | ||
| "microbatch_wise_load_balance_loss" | ||
| if self.top_k is None | ||
| else "unbiased_top_k_load_balance_loss" | ||
| ) | ||
|
|
||
| def _init_self_buffers(self, *, buffer_device: torch.device | None = None): | ||
| """Materialize the independent counter on the loss accumulator's device.""" | ||
| super()._init_self_buffers(buffer_device=buffer_device) | ||
| self.arrival_counter = torch.zeros( | ||
| 1, dtype=torch.int32, device=self.instance_acc.device | ||
| ) | ||
|
|
||
| def _compute_local_stats(self, scores_TE, routing_map_TE, padding_mask_T): | ||
| """Retain the configured normalization when the geometry is not fused.""" | ||
| if self.top_k is not None: | ||
| return UnbiasedTopKLoadBalanceLoss._compute_local_stats( | ||
| self, scores_TE, routing_map_TE, padding_mask_T | ||
| ) | ||
| return super()._compute_local_stats(scores_TE, routing_map_TE, padding_mask_T) |
There was a problem hiding this comment.
Design nit: one class with two modes. A single class chooses between two objectives based on top_k is None. The fallback reaches into a class it doesn't inherit from: UnbiasedTopKLoadBalanceLoss._compute_local_stats(self, ...) is an unbound call on a non-subclass instance. metric_name and the forward fallback also branch on top_k, which goes against the PR's own goal of keeping provided-map and unbiased objectives distinct.
A cleaner split: a shared fused mixin plus
FusedMicrobatchWiseLoadBalanceLoss(MicrobatchWiseLoadBalanceLoss)FusedUnbiasedTopKLoadBalanceLoss(UnbiasedTopKLoadBalanceLoss)
Each would inherit its native _compute_local_stats/forward fallback, so the unbound call and the top_k: int | None sentinel go away. metric_name would still need an override to keep the native metric keys. Not blocking, since this is recipe code.
| def __post_init__(self): | ||
| self.sharding_config = _with_local_counter_sharding(self.sharding_config) |
There was a problem hiding this comment.
Dead code in practice. Aux-loss configs are built without a sharding_config (config_utils.py:411, deepseek_v3.py:91). So _with_local_counter_sharding gets None, returns None, and arrival_counter is never declared V. Your own GPU test confirms it: test_spmd_loss_is_dp_local_and_carrier_keeps_its_sharding asserts that every buffer, including arrival_counter, is spmd.R. This matches how instance_acc is handled today, so it's probably fine. Either drop this hook, or, if the V declaration matters, make sure it actually takes effect and test for it.
| @unittest.skipUnless( | ||
| torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 3), | ||
| "GB300 specialization", | ||
| ) | ||
| class TestFusedDSv3SeqwiseLossGPU(unittest.TestCase): |
There was a problem hiding this comment.
CI coverage gap. This class only runs on sm_103 (GB300), but the PR is labeled ciflow/b200 (sm_100), so none of these tests run in CI. Several of them don't depend on the tuned CuTe kernel:
test_module_injection_and_checkpoint_metric(remat metric accounting)test_spmd_loss_is_dp_local_and_carrier_keeps_its_shardingtest_router_composition_preserves_model_and_checkpoint_statetest_router_composition_compiles_without_graph_breakstest_fullgraph_compile_and_cuda_graph
On B200, 4096x256 inputs take the generic Triton path, which also claims bitwise parity with native, so these tests would still be meaningful there. Suggest gating the class on capability[0] >= 8 (the same gate _supported uses) and keeping the (10, 3) skip only on tests that specifically check the tuned path (counter wrap, independent counters). As written, B200 CI runs only the parametrized generic test and test_general_sequence_module_compiles_and_replays.
| def _column_reduction_plan(tokens, experts, device): | ||
| """Derive ATen's FP32 output-vector, warp and CTA reduction geometry.""" | ||
| vector = 4 if experts % 4 == 0 else (2 if experts % 2 == 0 else 1) | ||
| threads = 512 // vector | ||
| output_power = min(threads, 1 << ((experts // vector).bit_length() - 1)) | ||
| initial_width = min(output_power, 32) | ||
| height = min(1 << (tokens.bit_length() - 1), threads // initial_width) | ||
| width = min(output_power, threads // height) | ||
| split_warps = tokens >= min(height * 16, 256) | ||
| reduction_threads = height if split_warps else 1 | ||
| per_thread = triton.cdiv(tokens, reduction_threads) | ||
| properties = torch.cuda.get_device_properties(device) | ||
| target_grid = properties.multi_processor_count * ( | ||
| properties.max_threads_per_multi_processor // (width * height) | ||
| ) | ||
| grid = triton.cdiv(experts, vector * width * (1 if split_warps else height)) | ||
| partitions = 1 | ||
| if split_warps and per_thread >= 256 and grid <= target_grid: | ||
| partitions = max( | ||
| triton.cdiv(per_thread, 256), | ||
| min(triton.cdiv(target_grid, grid), triton.cdiv(per_thread, 16)), | ||
| ) | ||
| return reduction_threads, partitions |
There was a problem hiding this comment.
Maintainability. This copies ATen's internal ReduceConfig heuristics (setReduceConfig in aten/src/ATen/native/cuda/Reduce.cuh). The bitwise-parity contract therefore depends on PyTorch internals that can change between releases, with no API guarantee. Please reference the exact upstream function and the PyTorch commit/version this mirrors, in the docstring or as a TODO. Then, when a parity test breaks after a PyTorch bump, it's obvious where to look. The literals 512, 16, 256 and the split_warps threshold need short comments pointing to the matching ATen constants.
Agent details (AI-generated)
Adapted from Yidi Wu (@ydwu4), #5259, revision
25dafa9. Preserve original attribution and
license; this change includes integration and generalization, not just a port.
Keep provided-map/L1 and unbiased top-k objectives distinct. Preserve native
coefficient/token normalization, gradient injection, metric accumulation and
DP-local training ownership. Padding contributes no counts, score sums or
gradients; unbiased selection is independent of biased/group-limited routing.
General fused CUDA kernels support positive T/E and configured top-k, while
the aligned production geometry keeps the tuned CuTe launch. Derive scratch
and native CTA partitioning from shapes and device resources. Explicitly fold
block.y partials in native order, independent of Triton layout; preserve
unaligned score-base offsets and scalar clamp-min's strict boundary derivative.
Temporary normalized scores/maps and partial buffers are not autograd saved
state. Save only original scores, frequencies and an optional padding mask.
The general implementation does not fall back to eager ATen for shape/tails.
Verification
Owning-commit CPU checks: 12 passed in 4.50s.
Strict relevant Sphinx builds and all non-type repository hooks passed.
Three unchanged main type findings remain attributed, not reported as a clean
full type run. External reviewer and CI acceptance are separate, still pending.
Complete-stack integration (not this PR in isolation): real four-GPU PP2/DP2/
EP2 and PP4/DP1/EP1 ran ten updates with exact model/optimizer tensor parity,
optimizer graphs and sequential validation. A fresh native PP2/PP4 checkpoint
matrix matched resumed state exactly. These are debug-model contract checks,
not 671B throughput or convergence results.
The final 207-case GPU matrix covers in-place gradients, dense attention, PP/FSDP train/eval graphs, shared-stream backward/capture and fused numerical contracts.
GPU JIT kernel/helper ASTs match the previously measured hot/general-shape
implementation. Prior matched GB300 microbenchmarks remain qualified evidence,
not a universal hardware or all-shape speedup claim. Original Yidi Wu source
credit remains above and in the implementation.
Removed unreachable eager wrapper branches and enforced tuned alignment; accepted unaligned scores still use the tested general fused kernels.