Skip to content

fix softmax collapse with improved ce perf #5661

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tianyu-l merged 1 commit into from
Nov 1, 2026
+389 −166
Merged

fix softmax collapse with improved ce perf #5661

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions tests/assets/losses/fake_pg/a10g/gpt_oss.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=4, TP=2, EP=4
# step loss grad_norm
1 7.639479637145996 159.70401000976562
2 6.5453200340271 32.60292053222656
3 7.166349411010742 7.431593894958496
4 6.35135555267334 4.453035354614258
5 5.490924835205078 3.7567286491394043
6 5.018106460571289 3.7990641593933105
7 5.918277263641357 39.103816986083984
8 5.388754844665527 10.378761291503906
9 5.015903472900391 4.206544876098633
10 5.0145368576049805 4.881128787994385
1 7.639479637145996 159.7058868408203
2 6.5406413078308105 32.58749771118164
3 7.167300224304199 7.427956581115723
4 6.353944301605225 4.440316200256348
5 5.492511749267578 3.7593331336975098
6 5.015379905700684 3.7947463989257812
7 5.938444137573242 40.8707275390625
8 5.39120626449585 10.596946716308594
9 5.017012596130371 4.2380900382995605
10 5.025738716125488 4.952903747558594
18 changes: 9 additions & 9 deletions tests/assets/losses/fake_pg/a10g/llama3.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@
# parallelism: FSDP=2, TP=2, CP=2
# step loss grad_norm
1 3.7498068809509277 879672.875
2 3.6602728366851807 394365.1875
3 3.6073100566864014 0.7835380434989929
4 3.636117458343506 0.7628054022789001
5 3.6130919456481934 883098.625
6 3.634310245513916 0.776549756526947
7 3.565253734588623 0.8194155693054199
8 3.5542614459991455 0.7671658992767334
9 3.476386547088623 996805.4375
10 3.413607120513916 2065079.625
2 3.6697301864624023 394365.1875
3 3.6074106693267822 0.7835395336151123
4 3.6361570358276367 0.7627957463264465
5 3.6227006912231445 884138.3125
6 3.634305000305176 0.776519238948822
7 3.5652427673339844 0.8193774819374084
8 3.5641703605651855 0.7671727538108826
9 3.486189842224121 990840.0
10 3.422697067260742 2058897.75
18 changes: 9 additions & 9 deletions tests/assets/losses/fake_pg/a10g/qwen3.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,12 @@
# parallelism: FSDP=2, TP=2, CP=2, EP=8
# step loss grad_norm
1 3.5141940116882324 74923716378624.0
2 3.4693727493286133 199539180961792.0
3 3.4716548919677734 3.3049397468566895
4 3.321932792663574 2.0839762687683105
5 3.496561288833618 2985306139656192.0
6 3.2210896015167236 2.095125436782837
7 3.126380681991577 1.9294486045837402
8 3.1302154064178467 1.576026439666748
9 2.987391710281372 1.85704301295829e+17
10 2.9555840492248535 2.4710699200544768e+17
2 3.481289863586426 199877359304704.0
3 3.4718329906463623 3.303375482559204
4 3.3228907585144043 2.084406614303589
5 3.5075387954711914 2971971071508480.0
6 3.2208504676818848 2.086042881011963
7 3.1266400814056396 1.9284083843231201
8 3.1411125659942627 1.5760431289672852
9 2.997920036315918 1.81068073078358e+17
10 2.9652247428894043 2.2956699840361267e+17
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/a10g/gpt_oss.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=4, TP=2, EP=4
# step loss grad_norm
1 7.955585479736328 46.24765396118164
2 5.301669120788574 15.41297721862793
3 5.297870635986328 4.662709712982178
4 4.499295711517334 2.6056394577026367
5 4.475003242492676 6.392116069793701
6 4.112522125244141 2.1256518363952637
7 4.056522369384766 2.215843915939331
8 3.9837710857391357 2.563532590866089
9 3.964216709136963 1.9628283977508545
10 3.858022689819336 1.8031314611434937
1 7.955585479736328 46.24832534790039
2 5.3015828132629395 15.409768104553223
3 5.297886848449707 4.665452480316162
4 4.500314712524414 2.604491949081421
5 4.483660697937012 6.5360212326049805
6 4.120448589324951 2.1533029079437256
7 4.064891815185547 2.2528274059295654
8 3.9945058822631836 2.5782649517059326
9 3.9672889709472656 1.9700688123703003
10 3.863626003265381 1.8391846418380737
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/a10g/llama3.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=2, TP=2, CP=2
# step loss grad_norm
1 8.023969650268555 1.5995080471038818
2 7.695952415466309 1.4733551740646362
3 6.94732666015625 1.877510905265808
4 6.103848457336426 2.368983268737793
5 6.189392566680908 1.7836463451385498
6 4.836900234222412 2.3710577487945557
7 4.538928985595703 2.295267105102539
8 4.793259620666504 1.82436203956604
9 4.086093425750732 1.9265673160552979
10 4.0527448654174805 1.8593209981918335
1 8.023969650268555 1.599506139755249
2 7.696089744567871 1.473367691040039
3 6.947298049926758 1.8775184154510498
4 6.103617191314697 2.369032859802246
5 6.189300537109375 1.7835813760757446
6 4.836762428283691 2.371192455291748
7 4.538995265960693 2.2955024242401123
8 4.793364524841309 1.824092984199524
9 4.086090087890625 1.9265425205230713
10 4.052403450012207 1.8593873977661133
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/a10g/llama3_pp.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=2, TP=2, PP=2
# step loss grad_norm
1 8.009735107421875 1.447790265083313
2 7.683969974517822 1.55353581905365
3 7.142031669616699 1.5556848049163818
4 6.181343078613281 2.1186180114746094
5 5.075255393981934 2.4584648609161377
6 4.709244728088379 2.239208221435547
7 4.471525192260742 1.979024887084961
8 4.387286186218262 1.85532546043396
9 4.382674217224121 1.673314094543457
10 4.196030139923096 1.692028284072876
1 8.009735107421875 1.4477896690368652
2 7.683945178985596 1.5535205602645874
3 7.142040252685547 1.555664300918579
4 6.181266784667969 2.1186206340789795
5 5.075375080108643 2.4584953784942627
6 4.7093424797058105 2.239180088043213
7 4.471562385559082 1.9791427850723267
8 4.387378692626953 1.8553677797317505
9 4.382554054260254 1.6732995510101318
10 4.196063995361328 1.6918935775756836
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/a10g/qwen3.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=2, TP=2, CP=2, EP=8
# step loss grad_norm
1 7.72227144241333 23.218080520629883
2 7.0880255699157715 20.9434871673584
3 6.549936294555664 3.7906463146209717
4 6.239657402038574 3.6388628482818604
5 6.505231857299805 15.887887001037598
6 5.856864929199219 3.538511037826538
7 5.697327613830566 3.4707326889038086
8 5.761327743530273 3.03719425201416
9 5.388110160827637 3.567819595336914
10 5.3545684814453125 3.4849300384521484
1 7.72227144241333 23.217422485351562
2 7.087889194488525 20.937583923339844
3 6.550065994262695 3.7905824184417725
4 6.23988151550293 3.638829231262207
5 6.5059123039245605 13.003586769104004
6 5.853584289550781 3.5368616580963135
7 5.6925272941589355 3.4690189361572266
8 5.758388519287109 3.03583025932312
9 5.383045196533203 3.565937042236328
10 5.34934663772583 3.4826102256774902
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/mi350x/gpt_oss.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=4, TP=2, EP=4
# step loss grad_norm
1 7.9550676345825195 46.24351501464844
2 5.300476551055908 15.368860244750977
3 5.305675506591797 4.6482014656066895
4 4.504476547241211 2.62524151802063
5 4.39337158203125 5.096515655517578
6 4.117705345153809 2.2558279037475586
7 4.047230243682861 2.1791443824768066
8 3.9648385047912598 2.4833266735076904
9 3.947808027267456 1.8302432298660278
10 3.846911907196045 1.6843359470367432
1 7.9550676345825195 46.2436637878418
2 5.30027961730957 15.363542556762695
3 5.303028583526611 4.639032363891602
4 4.502612113952637 2.6170601844787598
5 4.388518333435059 5.055764675140381
6 4.116694927215576 2.2071192264556885
7 4.0553483963012695 2.227717161178589
8 3.9775290489196777 2.5561912059783936
9 3.9506473541259766 1.8586903810501099
10 3.847172737121582 1.704336404800415
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/mi350x/llama3.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=2, TP=2, CP=2
# step loss grad_norm
1 8.02390193939209 1.5995241403579712
2 7.696142196655273 1.4733582735061646
3 6.947509765625 1.8774969577789307
4 6.103885650634766 2.368839979171753
5 6.189201831817627 1.7835946083068848
6 4.837381362915039 2.371023654937744
7 4.539409637451172 2.2957608699798584
8 4.793272018432617 1.824829339981079
9 4.0862298011779785 1.9266570806503296
10 4.052866458892822 1.8584532737731934
1 8.02390193939209 1.5995243787765503
2 7.696120262145996 1.4733744859695435
3 6.947292804718018 1.8774689435958862
4 6.103675842285156 2.368863821029663
5 6.189241409301758 1.7835688591003418
6 4.837277412414551 2.370814085006714
7 4.539156436920166 2.2959346771240234
8 4.793213844299316 1.825095772743225
9 4.086241722106934 1.9262988567352295
10 4.0526862144470215 1.8586567640304565
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/mi350x/llama3_pp.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=2, TP=2, PP=2
# step loss grad_norm
1 8.009737014770508 1.447791337966919
2 7.683963775634766 1.5535626411437988
3 7.141963958740234 1.5557489395141602
4 6.181270599365234 2.118598699569702
5 5.0750732421875 2.458409309387207
6 4.7094526290893555 2.23930287361145
7 4.47149133682251 1.9793468713760376
8 4.3872389793396 1.855665683746338
9 4.382683753967285 1.6734408140182495
10 4.19615364074707 1.6920582056045532
1 8.009737014770508 1.4477909803390503
2 7.683986663818359 1.5535576343536377
3 7.1419830322265625 1.5557408332824707
4 6.18121337890625 2.1186060905456543
5 5.075322151184082 2.4585039615631104
6 4.7094831466674805 2.23927903175354
7 4.4716477394104 1.9792524576187134
8 4.3872809410095215 1.8556113243103027
9 4.382562637329102 1.673511266708374
10 4.196207046508789 1.6919211149215698
20 changes: 10 additions & 10 deletions tests/assets/losses/real_pg/mi350x/qwen3.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,13 @@
# ngpu: 8
# parallelism: FSDP=2, TP=2, CP=2, EP=8
# step loss grad_norm
1 7.722476959228516 23.21984100341797
2 7.087940692901611 20.940082550048828
3 6.549875736236572 3.790707588195801
4 6.239684104919434 3.6387741565704346
5 6.504705905914307 16.865272521972656
6 5.858005523681641 3.5390546321868896
7 5.698677062988281 3.4708900451660156
8 5.762312412261963 3.0375165939331055
9 5.389161109924316 3.568202495574951
10 5.3555474281311035 3.485375165939331
1 7.722476959228516 23.219192504882812
2 7.087913513183594 20.92938232421875
3 6.549829483032227 3.790701150894165
4 6.239555358886719 3.6387856006622314
5 6.505892753601074 13.002161026000977
6 5.853765487670898 3.5368473529815674
7 5.6925506591796875 3.4689786434173584
8 5.758376598358154 3.0357329845428467
9 5.383083820343018 3.5659635066986084
10 5.349068641662598 3.4823544025421143
Loading
Loading