Repository navigation
Support explicit DSA top-k backends and tie policies - #108
Conversation
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
|
Reviewed at head 8918ef3. High-quality implementation — I verified the autograd input-count compatibility ( 1. [Risk — medium] External top-k path drops the The external early-returns in 2. [Risk — low/medium] The
3. [Risk — low] External top-k + indexer loss silently degrades to the O(sq·sk) dense reference path
4. [Nit] Inconsistent predicates in L659 gates the offset on 5. [Nit] Naive-path index dtype flips between int64 and int32 depending on config
Test coverage is strong (tie policies, legacy-cuDNN fallback, TP-local slice fallback, skip-layer reuse, backward compaction skip, monkeypatched "internal kernel must not run" assertions). The one assumption not defended across cuDNN versions is finding 1. |
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
Re-based directly on miles-main (no longer stacked on GDN2, which was closed as unadopted by any open-source model). Same content as the original port, re-derived from a fresh cherry-pick of upstream 788e9e1 so this PR's diff only carries KDA, not GDN2. Adds experimental_attention_variant="kda": a channel-wise Gated DeltaNet variant with direct Q/K/V/F/G projections and a separate beta projection, backed by fla.ops.kda.chunk_kda. Also brings in the generic primitives this commit packages alongside KDA: MLA output gating (attention_output_gate / gated_attention_proj_granularity), a HybridModel-aware qk_clip traversal, and a MoE router metric-indexing refactor. kda.py itself needed no hand-adaptation: it was already authored upstream against the same headwise/chunkwise-CP, selective-recompute base class our common.py/gdn.py carry, so it applies as a clean add. Kept the earlier decision not to eagerly normalize experimental_attention_variant from "gated_delta_net" to "gdn" in __post_init__ (re-introduced by this upstream commit): our fork's dsv4_hybrid/gated_delta_net validation checks still key off the literal "gated_delta_net" string, so mutating it before those checks run would silently disable them. is_gated_delta_net_variant (non-mutating, canonically defined in transformer_config.py per #102/#108) is used everywhere a GDN/KDA family check is needed instead. Adopted upstream's `not self.is_hybrid_model` condition in place of a narrower "only gated_delta_net requires linear_attention_freq" check for the same requirement, since it is the more general form already present in this file. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UawmJvYe7WkXApRw26ZFMd
2 participants
What does this PR do?
@HumansAnd
Configure ordinary DSA top-k independently of score/attention kernels.
dsa_indexer_topk_backend=None|torch|flashinfer,dsa_indexer_topk_deterministic=False,dsa_indexer_topk_tie_break=0|1|2(unset/small/large), with CLI flags.Noneretains kernel selection; reject CSA/V4 and non-FlashInfer tie options.top_k(sorted=False, deterministic=..., tie_break=..., dsa_graph_safe=True). Preserve scores/order; int32 indices;-inf→-1. Ties use original columns without score bias.dsa_kernel_backend=noneorattention_backend=unfusedenables reference training with selector-preserving backward.Issue tracking
2eb171ffe700d320ec0e74c7020c8094dc939668.radixark/Megatron-LM:miles-main, basea84b105473eca9c9c6e49dbf1c8aa460c5cf796e.Validation
Current CPU validation
8c1a1300c01577c4bb4fddcfaf70ba0cd28a83ee: merge scorer dispatch, remove gradient-tuple slicing and require selector fields directly; update fixtures, retain test cases and numerical paths.macOS/Torch 2.14.0 CPU doubles; overlapping counts. Python/diff/changed-file Black pass; no new Ruff findings, baseline style issues retained. No GPU/distributed/perf rerun; historical commits remain below.
Complete CPU output for this checkout:
Historical GPU validation at
524ff28fradixark/miles:dev-202609291651, digestsha256:1f5e2f9ea617010024d814fba9f48dba12fd90447a2a042c26d287ee422db138; driver 590.48.01; CUDA 13.0.88; torch 2.13.0+cu130; TE 2.17.0; cuDNN frontend 1.28.0; FlashInfer 0.6.18./root/dsa-native-glm52/megatronin that image; use installed DSA/FlashInfer and TE checkout. Tests set policies directly, without SGLang environment variables./opt/datadownload failures retained in successful raw output.Complete final focused GPU output (52 cases)
Complete final existing-default regression output
Paired TP1/CP4/EP4 trainer comparison
Both frozen-trajectory arms exited 0 on the same C2 B300 node/image above, each with four actor GPUs: TP1/CP4/EP4/PP1, expertTP1, no sequence parallelism. The separate eight-GPU rollout/training E2E follows.
3439ec75132d99dab443402e894eaa314e666d92a84b105473eca9c9c6e49dbf1c8aa460c5cf796e0d6cb7bee4ee1586ee64ba1fd3827a432dcbe657524ff28f2463b93c7c3ede29c7981cef499f7b2dPinaster/GLM-5.2_5layer(3 dense/2 MoE), top-k2048; 2 steps×64 samples×100 active response tokens. Identical frozen tokens/logprobs/rewards/masks/routing, fresh processes/optimizers, separate conversions of one BF16 HF checkpoint. Config/tokenizer/rollout hashes checked; weight payloads un-hashed.--allgather-cp, native zigzag--cp-comm-type allgather.actor_trainincludes optimizer;trainadds scoring/advantages/logging. Primary-rank wall timer (TP0,last PP,effective DP/CP0), no CUDA sync/events, not rank-max/kernel time. Step0 cold first-use/compile with persistent disk caches; step1 one observation. Excludes loading/conversion/launch/generation/publication/train_wait; W&B off.Negative time change means shorter execution.
Single subsequent step: actor 10.12% less time, total trainer 9.18% less, actor scoring 29.14% more, reference scoring 46.25% more; no steady-state throughput claim.
Frozen-rollout file hashes also match.
Train–inference KL: logged sample mean
train/train_rollout_kl, inference-sampled low-variance k3 estimate ofKL(inference || trainer)with FP32 clamps, not full-vocabulary KL. CP reduction differs by a few1e-9; reference KL and paired discrepancy are separate.train/train_rollout_kl)d=train_logp-rollout_logp,z=clamp(d,-20,20),mean(clamp(exp(z)-1-z,-10,10)), computed in FP32. Paired k3:delta=native_logp-Miles_logp,mean(exp(delta)-1-delta), unclamped. Samples come from inference, not the baseline trainer, so paired k3 is a discrepancy, not unbiased/exact policy KL.kl_lossis reference-policy andppo_kloptimization-policy diagnostics.Complete raw recorded paired perf, rollout, and training-step metrics
Restore each key's
prefix, then mergecommonwithdict(zip(columns,values)). All records/keys/values/arms/steps round-trip to the original raw evidence, which retains timestamps/PIDs.[ {"kind":"perf","prefix":"perf/","common":{"mfu_peak_tflops":2250.0},"columns":["update_weights_time","data_preprocess_time","train_wait_time","ref_log_probs_time","log_probs_time","actor_train_time","train_time","log_probs_tflops","ref_log_probs_tflops","actor_train_tflops","actor_train_tok_per_s","actor_train_mfu","step_time","wait_time_ratio"],"rows":[ {"record":0,"arm":"Miles","step":0,"values":[9.632110595703125e-05,0.006532907485961914,3.35579252243042,39.506826639175415,6.017252683639526,106.78244948387146,152.9441637992859,3.8693686599096466,0.5893403984450714,0.6541234743734369,153.43345352341498,0.00029072154416597196,156.2999563217163,0.0214702076789011]}, {"record":1,"arm":"Miles","step":1,"values":[0.00010251998901367188,0.009912490844726562,1.5883934497833252,0.26657557487487793,0.21686410903930664,30.52073574066162,31.644895553588867,105.10535757057274,85.50513202681296,2.24046824281824,525.5443425838033,0.0009957636634747734,33.23328900337219,0.047795252814795144]}, {"record":6,"arm":"native","step":0,"values":[0.0001010894775390625,0.009304046630859375,0.324535608291626,26.206274032592773,3.6332199573516846,78.26344037055969,108.74045443534851,6.4083565614352045,0.8884501827262794,0.8924844924753784,209.3442343248069,0.00039665977443350154,109.06499004364014,0.0029756167232195194]}, {"record":7,"arm":"native","step":1,"values":[0.00010442733764648438,0.009916543960571289,1.4353315830230713,0.38987064361572266,0.28004932403564453,27.433070182800293,28.739303588867188,81.3913042043224,58.464467889679256,2.492638947035271,584.6957665736078,0.001107839532015676,30.17463517189026,0.04756748755524909]} ]}, {"kind":"rollout","prefix":"rollout/","common":{"response_lengths":100.0,"rewards":0.0,"truncated":1.0,"raw_reward":0.0,"advantages":0.0,"returns":0.0},"columns":["rollout_log_probs","total_lengths","ref_log_probs","log_probs"],"rows":[ {"record":2,"arm":"Miles","step":0,"values":[-10.37608551979065,256.0,-10.376765310764313,-10.376266777515411]}, {"record":3,"arm":"Miles","step":1,"values":[-10.290511727333069,250.625,-10.290502071380615,-10.290469765663147]}, {"record":8,"arm":"native","step":0,"values":[-10.37608551979065,256.0,-10.376913666725159,-10.376191258430481]}, {"record":9,"arm":"native","step":1,"values":[-10.290511727333069,250.625,-10.29033637046814,-10.29027271270752]} ]}, {"kind":"step","prefix":"train/","common":{"loss":0.0,"pg_loss":0.0,"entropy_loss":0.0,"pg_clipfrac":0.0,"tis_clipfrac":0.0,"grad_norm":0.0,"lr-pg_0":1e-06,"lr-pg_1":1e-06,"lr-pg_2":1e-06},"columns":["ppo_kl","ess_ratio","train_rollout_logprob_abs_diff","train_rollout_kl","kl_loss","ois","tis","tis_abs","step"],"rows":[ {"record":4,"arm":"Miles","step":0,"values":[-0.0002223100746050477,0.9998570084571838,0.010888094082474709,9.71974222920835e-05,0.00017318082973361015,1.0002942085266113,0.9999154210090637,0.010886028409004211,0]}, {"record":5,"arm":"Miles","step":1,"values":[0.0001500502257840708,0.999860405921936,0.010912524536252022,9.61235273280181e-05,0.00015038735000416636,0.9999205470085144,1.000138282775879,0.010914952494204044,1]}, {"record":10,"arm":"native","step":0,"values":[0.0,1.0,0.011153023689985275,0.00010137323988601565,0.0001401206973241642,0.9999999403953552,0.999994695186615,0.011153263971209526,0]}, {"record":11,"arm":"native","step":1,"values":[0.0,1.0,0.010872159153223038,9.742614929564297e-05,0.0001125984126701951,1.000000238418579,1.0003365278244019,0.010876685380935669,1]} ]} ]$REPRO/public-repro, use the paired commits above, prepare seeds, then run GPU commands sequentially:Final full CP4 E2E
Completed
dsa-w4a16-cp4-finalon Miles41975f2d9589c417783ed38effc75865531fd2daand Megatron524ff28f2463b93c7c3ede29c7981cef499f7b2d: four actor GPUs (TP1/CP4/EP4/PP1/expertTP1), four rollout GPUs, two fresh steps, success2026-09-29 22:07:58, wrapper wall time726.1579617420211 s. Fresh trajectories are separate from the frozen pair above. Miles final adds only the early negative-loss guard and CPU test to paired0d6cb7be.The wrapper reused assets and called the real test's
execute(), overriding only run ID, independent Megatron import path, disabled W&B and launcher-owned Ray lifetime. Two successful publications:default → 1before rollout 0, then1 → 2before rollout 1. All 64 samples/rollout have the expected version; mixed-version ratios are zero. The unchanged loop skips final publication when no next rollout/evaluation is due.Eight CP shards reassemble into 64 samples/6,400 active tokens per step. Four dump/log checks pass
atol=1e-8, rtol=1e-4, maximum difference3.55270458385e-09; zero same-run paired deltas only audit extraction. All 128 rewards/losses/gradients are zero; responses hit 100 tokens. This verifies execution/version progress, not learning, quality or long-context boundaries. Numeric CI checkers remain disabled; this section makes no full-model validation claim.Raw final-E2E records (separate from paired replay):
Use the scalar encoding above; these records are from the separate final E2E.
[ {"kind":"perf","prefix":"perf/","common":{"mfu_peak_tflops":2250.0},"columns":["update_weights_implementation_time","finalize_and_resume_engines_time","update_weights_time","data_preprocess_time","train_wait_time","ref_log_probs_time","log_probs_time","actor_train_time","train_time","log_probs_tflops","ref_log_probs_tflops","actor_train_tflops","actor_train_tok_per_s","actor_train_mfu","step_time","wait_time_ratio"],"rows":[ {"record":0,"arm":"final-full-E2E","step":0,"values":[2.2860162258148193,0.023979663848876953,2.3967437744140625,0.009305000305175781,9.883866310119629,59.397878885269165,2.823371171951294,70.34786534309387,133.2052984237671,8.246513665697249,0.3919831716180397,0.9929072690100206,232.89974642575896,0.00044129211956000917,143.08916473388672,0.06907487599428909]}, {"record":1,"arm":"final-full-E2E","step":1,"values":[1.1596965789794922,0.024682998657226562,1.209233283996582,0.008936166763305664,3.3599729537963867,0.3271448612213135,0.2734978199005127,29.032727479934692,30.269611358642578,83.3409923819187,69.6742710238696,2.355298489322431,552.4799559767741,0.001046799328587747,33.629584312438965,0.09991122466992833]} ]}, {"kind":"rollout","prefix":"rollout/","common":{"response_lengths":100.0,"rewards":0.0,"truncated":1.0,"raw_reward":0.0,"advantages":0.0,"returns":0.0},"columns":["rollout_log_probs","total_lengths","ref_log_probs","log_probs"],"rows":[ {"record":2,"arm":"final-full-E2E","step":0,"values":[-10.322879791259766,256.0,-10.322923600673676,-10.322820544242859]}, {"record":3,"arm":"final-full-E2E","step":1,"values":[-10.288553476333618,250.625,-10.288492321968079,-10.2886803150177]} ]}, {"kind":"rollout_perf","prefix":"","common":{"rollout/raw_reward_unfiltered":0.0,"rollout/num_training_samples":64,"rollout/episode_raw_reward":0.0,"rollout/episode_response_length/mean":100.0,"rollout/episode_response_length/median":100.0,"rollout/episode_response_length/max":100,"rollout/episode_response_length/min":100,"rollout/episode_total_response_length/mean":100.0,"rollout/response_len/mean":100.0,"rollout/response_len/median":100.0,"rollout/response_len/max":100,"rollout/response_len/min":100,"rollout/zero_std/count_0":8,"rollout/zero_std/all_zero_percentage":0.0,"rollout/zero_std/all_one_percentage":0.0,"rollout/repetition_frac":0.0,"rollout/truncated_ratio":1.0,"rollout/weight_version/mixed_version_ratio":0.0},"columns":["rollout/prefix_cache_hit_rate","rollout/avg_cached_tokens_per_sample","rollout/weight_version/mean","rollout/weight_version/median","rollout/weight_version/max","rollout/weight_version/min","perf/rollout_time","perf/tokens_per_gpu_per_sec","perf/longest_sample_tokens_per_sec","perf/effective_tokens_per_gpu_per_sec","perf/longest_effective_sample_tokens_per_sec"],"rows":[ {"record":4,"arm":"final-full-E2E","step":0,"values":[0.05128205128205128,8.0,1.0,1.0,1,1,7.243313789367676,220.89337098009062,13.805835686255664,220.89337098009062,13.805835686255664]}, {"record":5,"arm":"final-full-E2E","step":1,"values":[0.006639004149377593,1.0,2.0,2.0,2,2,0.6896185874938965,2320.123078199601,145.00769238747506,2320.123078199601,145.00769238747506]} ]}, {"kind":"step","prefix":"train/","common":{"loss":0.0,"pg_loss":0.0,"entropy_loss":0.0,"pg_clipfrac":0.0,"ppo_kl":0.0,"ess_ratio":1.0,"tis_clipfrac":0.0,"grad_norm":0.0,"lr-pg_0":1e-06,"lr-pg_1":1e-06,"lr-pg_2":1e-06},"columns":["train_rollout_logprob_abs_diff","train_rollout_kl","kl_loss","ois","tis","tis_abs","step"],"rows":[ {"record":6,"arm":"final-full-E2E","step":0,"values":[0.011127699166536331,9.903085447149351e-05,0.00011141906725242734,0.9999999403953552,1.0001581907272339,0.011130120605230331,0]}, {"record":7,"arm":"final-full-E2E","step":1,"values":[0.010931901633739471,9.720325033413246e-05,0.00012288187281228602,1.000000238418579,0.9999712109565735,0.010932303965091705,1]} ]} ]Formatting and static checks
Historical scoped formatting at
524ff28f2463b93c7c3ede29c7981cef499f7b2d: six Python files versus fork base.tools/autoformat.shhardcodes NVIDIA's remote/base, so its commands ran with fork-local scope. Black 26.3.0, isort 5.13.2, Ruff 0.9.10, pylint 3.2.6.uv run --active --no-project isort/isort --checkruff check --no-fixgit diff --check/python -m py_compileblack --check --difftransformer_config.pyformatting changespylinttransformer_config.pylong lines; same diagnostics reproduced against unmodified fork basemypy --explicit-package-bases --follow-imports=skipContribution process
Pre-checks
tools/autoformat.shon my PR (equivalent fork-scoped commands run; existing baseline failures documented above)Code review
Ready for review.
128K (131,072 tokens) TP1/CP4/EP4 W4A16 trainer comparison
Both backends replay identical frozen five-layer GLM-5.2 trajectories; CP reconstruction verifies tokens, masks and inference logprobs.
Environment and provenance
{"digest":"sha256:686518570f7f460e73986d613457604e7012e2581f9f3d39568bc221ec7c277c","driver":"590.48.01","environment_repair":"Pinned opentelemetry-api==1.44.0 to match installed SDK 1.44.0; original image API 1.45.0 failed Ray dashboard startup before model execution.","hardware":"1 C2 node, 8 B300 (4 trainer GPUs, TP1/CP4/EP4)","hf_revision":"1c749139f70e158e4420ba67f342bef1de2e650d","image":"radixark/miles:dev-202609292119","megatron":"524ff28f2463b93c7c3ede29c7981cef499f7b2d","miles":"41975f2d9589c417783ed38effc75865531fd2da","packages":{"flashinfer-python":"0.6.18","miles":"0.1.1","nvidia-cudnn-frontend":"1.28.0","ray":"2.58.0","sglang":"0.5.21.dev65+g57e7d0e","tilelang":"0.1.14","torch":"2.13.0+cu130","transformer-engine":"2.17.0","transformers":"5.12.1"},"prefill_graphs":"Explicitly disabled (--sglang-cuda-graph-backend-prefill disabled); decode graph settings unchanged.","sglang":"57e7d0e848d8812c9deb4673b96e8e4626d6e794"}Both arms used the identical historical commits above.
perf/log_probs_timeperf/ref_log_probs_timeperf/actor_train_timeperf/train_timeperf/log_probs_timeperf/ref_log_probs_timeperf/actor_train_timeperf/train_timeThe timer definitions/exclusions above apply: step 0 includes first use; step 1 is one subsequent observation. Negative time change means shorter execution.
The sample-weighted train–inference KL and paired k3 definitions above apply. Dump scores are pre-update; paired k3 is a discrepancy, not full-vocabulary or unbiased policy KL. All 8 CP dump/log checks pass
atol=1e-8, rtol=1e-4; max absolute difference1.1017546056e-08.Complete raw scalar metrics (losslessly factored)
For each row, restore the key
prefixand mergecommonwithdict(zip(columns, values)). This retains every original scalar metric and is automatically round-trip checked.[ {"kind":"perf","prefix":"perf/","common":{"mfu_peak_tflops":2250.0},"columns":["update_weights_time","data_preprocess_time","train_wait_time","ref_log_probs_time","log_probs_time","actor_train_time","train_time","log_probs_tflops","ref_log_probs_tflops","actor_train_tflops","actor_train_tok_per_s","actor_train_mfu","step_time","wait_time_ratio"],"rows":[ {"record":0,"arm":"baseline","step":0,"values":[9.5367431640625e-05,0.5797655582427979,23.90319275856018,183.82437682151794,117.38612198829651,563.100072145462,865.0427393913269,484.4558257209162,309.36261903518954,302.9748714434151,14897.188643643123,0.1346554984192956,888.9459321498871,0.02688936626410009]}, {"record":1,"arm":"baseline","step":1,"values":[9.322166442871094e-05,1.0143084526062012,22.524129629135132,110.6300756931305,110.62006950378418,499.50444531440735,721.4849011898041,514.0874609021195,514.0409630899999,341.54885620820386,16793.860552572034,0.15179949164809062,744.0090308189392,0.03027400030930077]}, {"record":6,"arm":"candidate","step":0,"values":[0.00012087821960449219,1.014634609222412,14.903969287872314,66.53362512588501,30.610384464263916,161.60927772521973,259.4956245422363,1857.813668508717,854.7315819394853,1055.6644666039836,51906.722918859545,0.46918420737954825,274.39959383010864,0.054314837277419356]}, {"record":7,"arm":"candidate","step":1,"values":[0.00010132789611816406,1.0053410530090332,22.89513325691223,29.007310390472412,28.289104461669922,119.95538187026978,177.9873948097229,2010.2577207090353,1960.4847843699108,1422.2385799459707,69931.06827897203,0.6321060355315425,200.88252806663513,0.11397274555066154]} ]}, {"kind":"rollout","prefix":"rollout/","common":{"response_lengths":100.0,"rewards":0.0,"truncated":1.0,"raw_reward":0.0,"total_lengths":131072.0,"advantages":0.0,"returns":0.0},"columns":["rollout_log_probs","ref_log_probs","log_probs"],"rows":[ {"record":2,"arm":"baseline","step":0,"values":[-10.301874160766602,-10.30439281463623,-10.304338455200195]}, {"record":3,"arm":"baseline","step":1,"values":[-10.308305740356445,-10.311294555664062,-10.310310363769531]}, {"record":8,"arm":"candidate","step":0,"values":[-10.301874160766602,-10.304998397827148,-10.30412483215332]}, {"record":9,"arm":"candidate","step":1,"values":[-10.308305740356445,-10.311031341552734,-10.310235977172852]} ]}, {"kind":"step","prefix":"train/","common":{"loss":0.0,"pg_loss":0.0,"entropy_loss":0.0,"pg_clipfrac":0.0,"tis_clipfrac":0.0,"grad_norm":0.0,"lr-pg_0":1e-06,"lr-pg_1":1e-06,"lr-pg_2":1e-06},"columns":["ppo_kl","ess_ratio","train_rollout_logprob_abs_diff","train_rollout_kl","kl_loss","ois","tis","tis_abs","step"],"rows":[ {"record":4,"arm":"baseline","step":0,"values":[3.830439527519047e-05,0.9998349547386169,0.06168312206864357,0.0032339124009013176,0.0006978258606977761,1.0000450611114502,1.0007715225219727,0.06168534979224205,0]}, {"record":5,"arm":"baseline","step":1,"values":[5.625905032502487e-05,0.9998306632041931,0.061381224542856216,0.003178331768140197,0.000672891503199935,1.0000290870666504,1.001174807548523,0.06147060543298721,1]}, {"record":10,"arm":"candidate","step":0,"values":[0.0,1.0,0.061617545783519745,0.0032244131434708834,0.0006542504415847361,1.0,1.0009747743606567,0.06163110211491585,0]}, {"record":11,"arm":"candidate","step":1,"values":[0.0,1.0,0.06096087396144867,0.0031331873033195734,0.0006417945842258632,1.0,1.001202940940857,0.06105007976293564,1]} ]} ]All 128 rewards and every loss/gradient norm are zero; this is execution/numerics evidence, not reward-driven learning. Weight decay can still update weights.