From 06562f393dba9e79a37d733ac3487fc89b26c858 Mon Sep 17 00:00:00 2001 From: Chun Fang Date: Fri, 10 Jul 2026 08:14:27 +0000 Subject: [PATCH 1/3] Use eager flex attention for determinism on ROCm Compiled ROCm flex attention (torch.compile(flex_attention)) on gfx950 is run-to-run unstable without explicit kernel_options and amplifies tiny (~1e-6) bf16 activation drift between FSDP8 all-gather and HSDP(4,2) replicate paths, causing step-1 loss mismatches that fail the exact --assert-equal parity check. Run flex attention eagerly (uncompiled) in --debug.deterministic mode on ROCm, which matches varlen FA parity exactly. CUDA keeps the compiled (max_autotune disabled) path unchanged. Refresh the ROCm llama3 reference loss file to match the eager-flex curve. Co-authored-by: Cursor --- tests/assets/losses/llama3_rocm_mi350x.txt | 200 ++++++++++----------- torchtitan/distributed/utils.py | 28 +-- 2 files changed, 118 insertions(+), 110 deletions(-) diff --git a/tests/assets/losses/llama3_rocm_mi350x.txt b/tests/assets/losses/llama3_rocm_mi350x.txt index 9c8689594c..8df7b0093e 100644 --- a/tests/assets/losses/llama3_rocm_mi350x.txt +++ b/tests/assets/losses/llama3_rocm_mi350x.txt @@ -1,100 +1,100 @@ -1 8.048327445983887 -2 7.7511491775512695 -3 7.032481670379639 -4 6.154104232788086 -5 5.193375587463379 -6 4.581057548522949 -7 4.174039840698242 -8 3.94850492477417 -9 3.8914096355438232 -10 3.65614914894104 -11 3.531560182571411 -12 3.408013343811035 -13 3.293339490890503 -14 3.2310104370117188 -15 3.1887032985687256 -16 3.1206564903259277 -17 3.021003007888794 -18 3.0503194332122803 -19 3.0538227558135986 -20 2.9775760173797607 -21 2.966926336288452 -22 3.010988235473633 -23 2.968798875808716 -24 2.9324653148651123 -25 2.911179780960083 -26 2.9226012229919434 -27 2.9021105766296387 -28 2.890856981277466 -29 2.9541642665863037 -30 2.8875749111175537 -31 2.8661625385284424 -32 2.8760147094726562 -33 2.8741350173950195 -34 2.8855152130126953 -35 2.871234655380249 -36 2.8446168899536133 -37 2.8499021530151367 -38 2.8139476776123047 -39 2.8105669021606445 -40 2.7671191692352295 -41 2.8396270275115967 -42 2.826798677444458 -43 2.79144549369812 -44 2.8076703548431396 -45 2.807429075241089 -46 2.8387398719787598 -47 2.858567714691162 -48 2.8117449283599854 -49 2.8402206897735596 -50 2.817715644836426 -51 2.8175201416015625 -52 2.7817206382751465 -53 2.7895617485046387 -54 2.7608649730682373 -55 2.7678987979888916 -56 2.778493881225586 -57 2.812225818634033 -58 2.781548500061035 -59 2.793610095977783 -60 2.787827730178833 -61 2.7580902576446533 -62 2.7649357318878174 -63 2.774937868118286 -64 2.7940354347229004 -65 2.7498083114624023 -66 2.775907516479492 -67 2.7525010108947754 -68 2.8094029426574707 -69 2.763075113296509 -70 2.8257927894592285 -71 2.7717387676239014 -72 2.7618136405944824 -73 2.7250499725341797 -74 2.7488436698913574 -75 2.773621082305908 -76 2.7601230144500732 -77 2.782832145690918 -78 2.7656335830688477 -79 2.751183032989502 -80 2.769622325897217 -81 2.750093460083008 -82 2.778364658355713 -83 2.771414279937744 -84 2.7603936195373535 -85 2.760734796524048 -86 2.802809238433838 -87 2.7754392623901367 -88 2.7338876724243164 -89 2.738309383392334 -90 2.7482380867004395 -91 2.7457528114318848 -92 2.7811965942382812 -93 2.7664952278137207 -94 2.7834484577178955 -95 2.763306140899658 -96 2.7609148025512695 -97 2.7779273986816406 -98 2.7334446907043457 -99 2.7794971466064453 -100 2.772202491760254 +1 8.048333168029785 +2 7.751149654388428 +3 7.032470226287842 +4 6.154115200042725 +5 5.193387985229492 +6 4.5810627937316895 +7 4.174030303955078 +8 3.9485042095184326 +9 3.8914008140563965 +10 3.6561331748962402 +11 3.5315587520599365 +12 3.408001184463501 +13 3.2933356761932373 +14 3.231022596359253 +15 3.188683032989502 +16 3.1206631660461426 +17 3.0210442543029785 +18 3.050333261489868 +19 3.0538251399993896 +20 2.977559804916382 +21 2.9669554233551025 +22 3.010943651199341 +23 2.968801975250244 +24 2.93245530128479 +25 2.9111685752868652 +26 2.922588348388672 +27 2.902114152908325 +28 2.890869379043579 +29 2.9541852474212646 +30 2.887592315673828 +31 2.866168737411499 +32 2.8760170936584473 +33 2.874119520187378 +34 2.885528564453125 +35 2.8712432384490967 +36 2.8446245193481445 +37 2.8499231338500977 +38 2.8139071464538574 +39 2.8105883598327637 +40 2.7671310901641846 +41 2.8396549224853516 +42 2.826801300048828 +43 2.791419744491577 +44 2.8076796531677246 +45 2.807339906692505 +46 2.838744878768921 +47 2.858563184738159 +48 2.8117125034332275 +49 2.8401975631713867 +50 2.8177144527435303 +51 2.8175175189971924 +52 2.78169846534729 +53 2.789565086364746 +54 2.760869264602661 +55 2.767879009246826 +56 2.778499126434326 +57 2.8122334480285645 +58 2.7815592288970947 +59 2.7935986518859863 +60 2.7877984046936035 +61 2.7581119537353516 +62 2.7649855613708496 +63 2.7749454975128174 +64 2.793992757797241 +65 2.749854564666748 +66 2.7759299278259277 +67 2.752507209777832 +68 2.809450149536133 +69 2.7630889415740967 +70 2.82578706741333 +71 2.771724224090576 +72 2.761836528778076 +73 2.725079298019409 +74 2.7488763332366943 +75 2.773552894592285 +76 2.7601332664489746 +77 2.7828969955444336 +78 2.765627861022949 +79 2.7512450218200684 +80 2.7696332931518555 +81 2.7501087188720703 +82 2.778395652770996 +83 2.771440029144287 +84 2.7603678703308105 +85 2.76071834564209 +86 2.8027772903442383 +87 2.775432586669922 +88 2.7338902950286865 +89 2.738295555114746 +90 2.7482383251190186 +91 2.7457945346832275 +92 2.7812249660491943 +93 2.7665035724639893 +94 2.7834959030151367 +95 2.7633213996887207 +96 2.760951280593872 +97 2.777953863143921 +98 2.7334725856781006 +99 2.779473066329956 +100 2.772204875946045 diff --git a/torchtitan/distributed/utils.py b/torchtitan/distributed/utils.py index 6a2018e0f4..6dfd0a26a9 100644 --- a/torchtitan/distributed/utils.py +++ b/torchtitan/distributed/utils.py @@ -189,20 +189,28 @@ def set_determinism( # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - # Ensure flex_attention is compiled without max-autotune. This is needed to ensure - # reproducibility, since the autotune results may not be deterministic. We disable - # autotune in-place on FlexAttention.inductor_configs (rather than recompiling with - # no options) so the regional-inductor scoop configs are preserved. from torch.nn.attention.flex_attention import flex_attention from torchtitan.models.common.attention import FlexAttention - FlexAttention.inductor_configs["max_autotune"] = False - FlexAttention.inductor_configs["coordinate_descent_tuning"] = False - # pyrefly: ignore [no-matching-overload] - FlexAttention._compiled_flex_attn = torch.compile( - flex_attention, options=FlexAttention.inductor_configs - ) + if torch.version.hip is not None: + # Compiled ROCm flex attention is not deterministic. + # Falling back to eager (non-compiled) flex_attention for determinism on ROCm. + logger.info( + "Using eager (non-compiled) flex_attention for determinism on ROCm." + ) + FlexAttention._compiled_flex_attn = flex_attention + else: + # Ensure flex_attention is compiled without max-autotune. This is needed to ensure + # reproducibility, since the autotune results may not be deterministic. We disable + # autotune in-place on FlexAttention.inductor_configs (rather than recompiling with + # no options) so the regional-inductor scoop configs are preserved. + FlexAttention.inductor_configs["max_autotune"] = False + FlexAttention.inductor_configs["coordinate_descent_tuning"] = False + # pyrefly: ignore [no-matching-overload] + FlexAttention._compiled_flex_attn = torch.compile( + flex_attention, options=FlexAttention.inductor_configs + ) if debug_config.detect_anomaly: logger.warning( From e0b4995f82ed78f46b12562bb31437154deeb1f8 Mon Sep 17 00:00:00 2001 From: Chun Fang Date: Fri, 10 Jul 2026 14:39:49 +0000 Subject: [PATCH 2/3] Regenerate ROCm llama3 reference losses in clean environment The previous reference was generated while another job contended for the GPUs, introducing ~1e-6 run-to-run numerical drift. Regenerated from a fresh FSDP8 eager-flex run on an uncontended node; two back-to-back runs now produce bitwise-identical losses, so the exact --assert-equal reference check reproduces reliably. Co-authored-by: Cursor --- tests/assets/losses/llama3_rocm_mi350x.txt | 198 ++++++++++----------- 1 file changed, 99 insertions(+), 99 deletions(-) diff --git a/tests/assets/losses/llama3_rocm_mi350x.txt b/tests/assets/losses/llama3_rocm_mi350x.txt index 8df7b0093e..5b52727466 100644 --- a/tests/assets/losses/llama3_rocm_mi350x.txt +++ b/tests/assets/losses/llama3_rocm_mi350x.txt @@ -1,100 +1,100 @@ 1 8.048333168029785 -2 7.751149654388428 -3 7.032470226287842 -4 6.154115200042725 -5 5.193387985229492 -6 4.5810627937316895 -7 4.174030303955078 -8 3.9485042095184326 -9 3.8914008140563965 -10 3.6561331748962402 -11 3.5315587520599365 -12 3.408001184463501 -13 3.2933356761932373 -14 3.231022596359253 -15 3.188683032989502 -16 3.1206631660461426 -17 3.0210442543029785 -18 3.050333261489868 -19 3.0538251399993896 -20 2.977559804916382 -21 2.9669554233551025 -22 3.010943651199341 -23 2.968801975250244 -24 2.93245530128479 -25 2.9111685752868652 -26 2.922588348388672 -27 2.902114152908325 -28 2.890869379043579 -29 2.9541852474212646 -30 2.887592315673828 -31 2.866168737411499 -32 2.8760170936584473 -33 2.874119520187378 -34 2.885528564453125 -35 2.8712432384490967 -36 2.8446245193481445 -37 2.8499231338500977 -38 2.8139071464538574 -39 2.8105883598327637 -40 2.7671310901641846 -41 2.8396549224853516 -42 2.826801300048828 -43 2.791419744491577 -44 2.8076796531677246 -45 2.807339906692505 -46 2.838744878768921 -47 2.858563184738159 -48 2.8117125034332275 -49 2.8401975631713867 -50 2.8177144527435303 -51 2.8175175189971924 -52 2.78169846534729 -53 2.789565086364746 -54 2.760869264602661 -55 2.767879009246826 -56 2.778499126434326 -57 2.8122334480285645 -58 2.7815592288970947 -59 2.7935986518859863 -60 2.7877984046936035 -61 2.7581119537353516 -62 2.7649855613708496 -63 2.7749454975128174 -64 2.793992757797241 -65 2.749854564666748 -66 2.7759299278259277 -67 2.752507209777832 -68 2.809450149536133 -69 2.7630889415740967 -70 2.82578706741333 -71 2.771724224090576 -72 2.761836528778076 -73 2.725079298019409 -74 2.7488763332366943 -75 2.773552894592285 -76 2.7601332664489746 -77 2.7828969955444336 -78 2.765627861022949 -79 2.7512450218200684 -80 2.7696332931518555 -81 2.7501087188720703 -82 2.778395652770996 -83 2.771440029144287 -84 2.7603678703308105 -85 2.76071834564209 -86 2.8027772903442383 -87 2.775432586669922 -88 2.7338902950286865 -89 2.738295555114746 -90 2.7482383251190186 -91 2.7457945346832275 -92 2.7812249660491943 -93 2.7665035724639893 -94 2.7834959030151367 -95 2.7633213996887207 -96 2.760951280593872 -97 2.777953863143921 -98 2.7334725856781006 -99 2.779473066329956 -100 2.772204875946045 +2 7.7511491775512695 +3 7.032463073730469 +4 6.154108047485352 +5 5.19338321685791 +6 4.581062316894531 +7 4.17402458190918 +8 3.9485011100769043 +9 3.8914105892181396 +10 3.6561455726623535 +11 3.531575918197632 +12 3.40800404548645 +13 3.293339729309082 +14 3.231020450592041 +15 3.18870210647583 +16 3.120664596557617 +17 3.0210518836975098 +18 3.0503437519073486 +19 3.053832530975342 +20 2.9775662422180176 +21 2.9669432640075684 +22 3.010946035385132 +23 2.9688079357147217 +24 2.9324467182159424 +25 2.91117787361145 +26 2.9225873947143555 +27 2.9020886421203613 +28 2.8908493518829346 +29 2.9541983604431152 +30 2.8875508308410645 +31 2.866196632385254 +32 2.8760130405426025 +33 2.8740994930267334 +34 2.8855278491973877 +35 2.8711819648742676 +36 2.8445587158203125 +37 2.849912405014038 +38 2.813920021057129 +39 2.8105406761169434 +40 2.7671115398406982 +41 2.839630603790283 +42 2.8267719745635986 +43 2.7914342880249023 +44 2.807629108428955 +45 2.807225227355957 +46 2.838663339614868 +47 2.8584837913513184 +48 2.811601400375366 +49 2.840069532394409 +50 2.817638635635376 +51 2.8174750804901123 +52 2.781676769256592 +53 2.78944730758667 +54 2.760826826095581 +55 2.7678024768829346 +56 2.778507709503174 +57 2.812169075012207 +58 2.781461238861084 +59 2.7935383319854736 +60 2.787777900695801 +61 2.758037567138672 +62 2.76495623588562 +63 2.7749221324920654 +64 2.794041633605957 +65 2.7497072219848633 +66 2.7758588790893555 +67 2.7524397373199463 +68 2.8093056678771973 +69 2.7630515098571777 +70 2.8258039951324463 +71 2.771604061126709 +72 2.7618227005004883 +73 2.7250170707702637 +74 2.7488081455230713 +75 2.773557186126709 +76 2.7600669860839844 +77 2.7827961444854736 +78 2.7656686305999756 +79 2.7511205673217773 +80 2.7695093154907227 +81 2.7500877380371094 +82 2.7782886028289795 +83 2.771364450454712 +84 2.7604129314422607 +85 2.7606608867645264 +86 2.8027162551879883 +87 2.7753896713256836 +88 2.733940601348877 +89 2.7383008003234863 +90 2.748220682144165 +91 2.745736837387085 +92 2.78118896484375 +93 2.7664875984191895 +94 2.7835144996643066 +95 2.763306140899658 +96 2.760876178741455 +97 2.77791428565979 +98 2.733462333679199 +99 2.7794244289398193 +100 2.772249937057495 From 5776ab9b36c831c2e2ec9afd99f57f6e7e7fc9c8 Mon Sep 17 00:00:00 2001 From: yuankaichen-amd Date: Sat, 11 Jul 2026 14:20:10 +0000 Subject: [PATCH 3/3] Regenerate qwen3 MoE ROCm reference losses with eager flex attention --- tests/assets/losses/qwen3_moe_rocm_mi350x.txt | 200 +++++++++--------- 1 file changed, 100 insertions(+), 100 deletions(-) diff --git a/tests/assets/losses/qwen3_moe_rocm_mi350x.txt b/tests/assets/losses/qwen3_moe_rocm_mi350x.txt index 91ec4bb50b..f07a45caf0 100644 --- a/tests/assets/losses/qwen3_moe_rocm_mi350x.txt +++ b/tests/assets/losses/qwen3_moe_rocm_mi350x.txt @@ -1,100 +1,100 @@ -1 8.039478302001953 -2 7.3455915451049805 -3 5.886664390563965 -4 5.181613445281982 -5 4.755587577819824 -6 4.551773548126221 -7 4.365748882293701 -8 4.220365524291992 -9 4.27623176574707 -10 4.035205841064453 -11 3.930511713027954 -12 3.890399217605591 -13 3.826479196548462 -14 3.697934150695801 -15 3.676894187927246 -16 3.6526780128479004 -17 3.4601187705993652 -18 3.407802104949951 -19 3.4914276599884033 -20 3.405426263809204 -21 3.3306427001953125 -22 3.320408582687378 -23 3.2711071968078613 -24 3.190694808959961 -25 3.179900884628296 -26 3.1389572620391846 -27 3.0819246768951416 -28 3.1187551021575928 -29 3.0865421295166016 -30 3.0367109775543213 -31 3.0023996829986572 -32 2.9462947845458984 -33 2.96773099899292 -34 2.9433414936065674 -35 2.985591173171997 -36 2.9773929119110107 -37 2.9675049781799316 -38 2.9271206855773926 -39 2.9594998359680176 -40 2.965256690979004 -41 3.0217385292053223 -42 2.9658241271972656 -43 2.8713490962982178 -44 2.964686870574951 -45 2.9258570671081543 -46 2.8722586631774902 -47 2.8407859802246094 -48 2.9174423217773438 -49 2.866563081741333 -50 2.8704583644866943 -51 2.8910083770751953 -52 2.8712692260742188 -53 2.862722873687744 -54 2.907217264175415 -55 2.8347325325012207 -56 2.8669302463531494 -57 2.870532989501953 -58 2.834479331970215 -59 2.82023286819458 -60 2.8041574954986572 -61 2.8596506118774414 -62 2.767400026321411 -63 2.7994165420532227 -64 2.8032584190368652 -65 2.8946614265441895 -66 2.8474836349487305 -67 2.8692901134490967 -68 2.8682971000671387 -69 2.8341708183288574 -70 2.8166182041168213 -71 2.7966384887695312 -72 2.785504102706909 -73 2.7746710777282715 -74 2.804699420928955 -75 2.793273448944092 -76 2.824176788330078 -77 2.7725069522857666 -78 2.8015503883361816 -79 2.766071081161499 -80 2.7475364208221436 -81 2.810965061187744 -82 2.8011093139648438 -83 2.847364664077759 -84 2.775377035140991 -85 2.800250768661499 -86 2.832399606704712 -87 2.8371798992156982 -88 2.7580952644348145 -89 2.8046157360076904 -90 2.8352174758911133 -91 2.7917966842651367 -92 2.8108761310577393 -93 2.8189456462860107 -94 2.899336576461792 -95 2.77498197555542 -96 2.86659574508667 -97 2.8083336353302 -98 2.740171194076538 -99 2.766096591949463 -100 2.7913966178894043 +1 8.039487838745117 +2 7.345444202423096 +3 5.8864569664001465 +4 5.181918621063232 +5 4.755563735961914 +6 4.552217483520508 +7 4.365880012512207 +8 4.22066593170166 +9 4.275698661804199 +10 4.035106658935547 +11 3.9316513538360596 +12 3.889565944671631 +13 3.826673746109009 +14 3.698694944381714 +15 3.6758856773376465 +16 3.65238881111145 +17 3.4600486755371094 +18 3.4068470001220703 +19 3.490159749984741 +20 3.4041874408721924 +21 3.3300230503082275 +22 3.3213939666748047 +23 3.269399881362915 +24 3.1864099502563477 +25 3.177971839904785 +26 3.1385791301727295 +27 3.081350803375244 +28 3.117417812347412 +29 3.0856497287750244 +30 3.038393974304199 +31 2.999509811401367 +32 2.9440619945526123 +33 2.9685757160186768 +34 2.9416675567626953 +35 2.984910011291504 +36 2.9781534671783447 +37 2.96781587600708 +38 2.9260241985321045 +39 2.9601118564605713 +40 2.965358257293701 +41 3.0216760635375977 +42 2.9644174575805664 +43 2.8688931465148926 +44 2.963170051574707 +45 2.925992012023926 +46 2.8718581199645996 +47 2.84104585647583 +48 2.9172792434692383 +49 2.8664212226867676 +50 2.8697421550750732 +51 2.8896594047546387 +52 2.8695273399353027 +53 2.8619370460510254 +54 2.9065823554992676 +55 2.8338732719421387 +56 2.8684566020965576 +57 2.870880126953125 +58 2.8348171710968018 +59 2.8202507495880127 +60 2.8038015365600586 +61 2.8593192100524902 +62 2.767629623413086 +63 2.7992656230926514 +64 2.803349018096924 +65 2.8964009284973145 +66 2.846656560897827 +67 2.871981620788574 +68 2.8704471588134766 +69 2.834507942199707 +70 2.8173625469207764 +71 2.793182373046875 +72 2.7870612144470215 +73 2.775667428970337 +74 2.804844379425049 +75 2.792123794555664 +76 2.8238449096679688 +77 2.7727913856506348 +78 2.8021910190582275 +79 2.76702880859375 +80 2.7479591369628906 +81 2.8113951683044434 +82 2.801135540008545 +83 2.845879554748535 +84 2.776027202606201 +85 2.7993133068084717 +86 2.8321409225463867 +87 2.8379364013671875 +88 2.7581987380981445 +89 2.80420184135437 +90 2.835993766784668 +91 2.7928967475891113 +92 2.8110811710357666 +93 2.820006847381592 +94 2.894343137741089 +95 2.77492094039917 +96 2.8666553497314453 +97 2.8066415786743164 +98 2.740886688232422 +99 2.7648468017578125 +100 2.790055513381958