[llama4] use grouped_mm in moe for sm90 #2755

IvanKobzarev · 2025-05-21T14:59:31Z

Stack from ghstack (oldest at bottom):

-> [llama4] use grouped_mm in moe for sm90 #2755

Enablement copy pasted from torchtitan.
There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty)

        out = _grouped_mm(h, w2, offs=offsets)
        out[offsets[-1] :].zero_()

scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index.
To prevent copying unitialized data - zero_() after grouped_mm for now.

grouped_mm in eager:
eager tps_avg: 3.2182 -> 4.44

grouped_mm in compile:
compile tps_avg: 4.33463 -> 8.70872

Observing NaN loss in compile :( with/without grouped_mm - will be debugging this separately

no_compile + no grouped_mm:
BASELINE:

Step 1 | loss:4.149928092956543 lr:2e-05 tokens_per_second_per_gpu:2.203409433364868 peak_memory_active:28.82542610168457 peak_memory_alloc:28.82542610168457 peak_memory_reserved:40.197265625
Step 2 | loss:2.3132553100585938 lr:2e-05 tokens_per_second_per_gpu:2.492661952972412 peak_memory_active:28.804787158966064 peak_memory_alloc:28.804787158966064 peak_memory_reserved:40.251953125
Step 3 | loss:1.6078141927719116 lr:2e-05 tokens_per_second_per_gpu:3.8224945068359375 peak_memory_active:28.853728771209717 peak_memory_alloc:28.853728771209717 peak_memory_reserved:40.251953125
Step 4 | loss:1.4519065618515015 lr:2e-05 tokens_per_second_per_gpu:3.001920700073242 peak_memory_active:28.84491729736328 peak_memory_alloc:28.84491729736328 peak_memory_reserved:40.251953125
Step 5 | loss:1.2131776809692383 lr:2e-05 tokens_per_second_per_gpu:3.3134357929229736 peak_memory_active:28.804935932159424 peak_memory_alloc:28.804935932159424 peak_memory_reserved:40.251953125
Step 6 | loss:1.411360263824463 lr:2e-05 tokens_per_second_per_gpu:4.279047966003418 peak_memory_active:28.83890438079834 peak_memory_alloc:28.83890438079834 peak_memory_reserved:40.251953125
Step 7 | loss:1.1743241548538208 lr:2e-05 tokens_per_second_per_gpu:3.129912853240967 peak_memory_active:28.785582065582275 peak_memory_alloc:28.785582065582275 peak_memory_reserved:40.251953125
Step 8 | loss:1.3950272798538208 lr:2e-05 tokens_per_second_per_gpu:3.4216036796569824 peak_memory_active:28.967872619628906 peak_memory_alloc:28.967872619628906 peak_memory_reserved:40.251953125
Step 9 | loss:1.2500101327896118 lr:2e-05 tokens_per_second_per_gpu:3.410867214202881 peak_memory_active:28.854501247406006 peak_memory_alloc:28.854501247406006 peak_memory_reserved:40.251953125
Step 10 | loss:1.1264036893844604 lr:2e-05 tokens_per_second_per_gpu:3.1065995693206787 peak_memory_active:28.782616138458252 peak_memory_alloc:28.782616138458252 peak_memory_reserved:40.251953125
 /tmp/torchtune/llama4_17Bx16E/full/logs
└─ $ tune_logs_tps log_1747910831.txt
tps_avg: 3.2182  n= 10
peak_memory_alloc max   28.97
peak_memory_reser max   40.25





no_compile + grouped_mm:
______

EAGER_GROUPED_MM:
└─ $ cat log_1747915295.txt
Step 1 | loss:4.163714408874512 lr:2e-05 tokens_per_second_per_gpu:2.171847343444824 peak_memory_active:27.444169998168945 peak_memory_alloc:27.444169998168945 peak_memory_reserved:38.9609375
Step 2 | loss:2.283238172531128 lr:2e-05 tokens_per_second_per_gpu:3.158111810684204 peak_memory_active:27.428916931152344 peak_memory_alloc:27.428916931152344 peak_memory_reserved:39.017578125
Step 3 | loss:1.5704680681228638 lr:2e-05 tokens_per_second_per_gpu:4.869690418243408 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125
Step 4 | loss:1.4485288858413696 lr:2e-05 tokens_per_second_per_gpu:4.026949405670166 peak_memory_active:27.46179723739624 peak_memory_alloc:27.46179723739624 peak_memory_reserved:39.017578125
Step 5 | loss:1.1918020248413086 lr:2e-05 tokens_per_second_per_gpu:4.502264976501465 peak_memory_active:27.428401947021484 peak_memory_alloc:27.428401947021484 peak_memory_reserved:39.017578125
Step 6 | loss:1.3454418182373047 lr:2e-05 tokens_per_second_per_gpu:5.665891170501709 peak_memory_active:27.456215858459473 peak_memory_alloc:27.456215858459473 peak_memory_reserved:39.017578125
Step 7 | loss:1.185951590538025 lr:2e-05 tokens_per_second_per_gpu:4.356126308441162 peak_memory_active:27.412775993347168 peak_memory_alloc:27.412775993347168 peak_memory_reserved:39.017578125
Step 8 | loss:1.3995944261550903 lr:2e-05 tokens_per_second_per_gpu:4.84677267074585 peak_memory_active:27.562548637390137 peak_memory_alloc:27.562548637390137 peak_memory_reserved:39.017578125
Step 9 | loss:1.2516542673110962 lr:2e-05 tokens_per_second_per_gpu:4.785239219665527 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125
Step 10 | loss:1.1245404481887817 lr:2e-05 tokens_per_second_per_gpu:4.080092906951904 peak_memory_active:27.41029644012451 peak_memory_alloc:27.41029644012451 peak_memory_reserved:39.017578125
Step 11 | loss:1.2505868673324585 lr:2e-05 tokens_per_second_per_gpu:4.96058988571167 peak_memory_active:27.416271209716797 peak_memory_alloc:27.416271209716797 peak_memory_reserved:39.0234375
Step 12 | loss:1.0976852178573608 lr:2e-05 tokens_per_second_per_gpu:4.555960655212402 peak_memory_active:27.415279388427734 peak_memory_alloc:27.415279388427734 peak_memory_reserved:39.0234375
Step 13 | loss:1.176978349685669 lr:2e-05 tokens_per_second_per_gpu:5.802102565765381 peak_memory_active:27.413253784179688 peak_memory_alloc:27.413253784179688 peak_memory_reserved:39.0234375
 /tmp/torchtune/llama4_17Bx16E/full/logs
└─ $ tune_logs_tps log_1747915295.txt
tps_avg: 4.44474  n= 13
peak_memory_alloc max   27.56
peak_memory_reser max   39.02








compile + no_grouped_mm:

COMPILED BASELINE:

└─ $ cat log_1747915995.txt
Step 1 | loss:4.1595354080200195 lr:2e-05 tokens_per_second_per_gpu:0.42956066131591797 peak_memory_active:28.81977891921997 peak_memory_alloc:28.81977891921997 peak_memory_reserved:40.236328125
Step 2 | loss:2.2778103351593018 lr:2e-05 tokens_per_second_per_gpu:2.819465160369873 peak_memory_active:28.79981756210327 peak_memory_alloc:28.79981756210327 peak_memory_reserved:40.2890625
Step 3 | loss:1.786258578300476 lr:2e-05 tokens_per_second_per_gpu:5.622256755828857 peak_memory_active:28.846065998077393 peak_memory_alloc:28.846065998077393 peak_memory_reserved:40.2890625
Step 4 | loss:1.600595474243164 lr:2e-05 tokens_per_second_per_gpu:2.996225357055664 peak_memory_active:28.837817668914795 peak_memory_alloc:28.837817668914795 peak_memory_reserved:40.2890625
Step 5 | loss:1.627750039100647 lr:2e-05 tokens_per_second_per_gpu:6.902056694030762 peak_memory_active:28.79964590072632 peak_memory_alloc:28.79964590072632 peak_memory_reserved:40.2890625
Step 6 | loss:1.6751512289047241 lr:2e-05 tokens_per_second_per_gpu:3.8652405738830566 peak_memory_active:28.833858966827393 peak_memory_alloc:28.833858966827393 peak_memory_reserved:40.2890625
Step 7 | loss:1.5953153371810913 lr:2e-05 tokens_per_second_per_gpu:2.8824896812438965 peak_memory_active:28.781311511993408 peak_memory_alloc:28.781311511993408 peak_memory_reserved:40.2890625
Step 8 | loss:1.6390446424484253 lr:2e-05 tokens_per_second_per_gpu:4.305164813995361 peak_memory_active:28.948315620422363 peak_memory_alloc:28.948315620422363 peak_memory_reserved:40.2890625
Step 9 | loss:1.53915536403656 lr:2e-05 tokens_per_second_per_gpu:4.096757888793945 peak_memory_active:28.84673547744751 peak_memory_alloc:28.84673547744751 peak_memory_reserved:40.2890625
Step 10 | loss:1.4715595245361328 lr:2e-05 tokens_per_second_per_gpu:6.167354106903076 peak_memory_active:28.778648853302002 peak_memory_alloc:28.778648853302002 peak_memory_reserved:40.2890625
Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.351371765136719 peak_memory_active:28.787389278411865 peak_memory_alloc:28.787389278411865 peak_memory_reserved:40.2890625
 /tmp/torchtune/llama4_17Bx16E/full/logs
└─ $ tune_logs_tps log_1747915995.txt
tps_avg: 4.33463  n= 12
peak_memory_alloc max   28.95
peak_memory_reser max   40.29




COMPILED + GROUPED_MM



└─ $ cat log_1747916676.txt
Step 1 | loss:4.163997650146484 lr:2e-05 tokens_per_second_per_gpu:0.9783019423484802 peak_memory_active:27.437111377716064 peak_memory_alloc:27.437111377716064 peak_memory_reserved:38.978515625
Step 2 | loss:2.2560360431671143 lr:2e-05 tokens_per_second_per_gpu:2.861786365509033 peak_memory_active:27.42238759994507 peak_memory_alloc:27.42238759994507 peak_memory_reserved:39.03515625
Step 3 | loss:1.8066692352294922 lr:2e-05 tokens_per_second_per_gpu:7.731247901916504 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625
Step 4 | loss:1.7610063552856445 lr:2e-05 tokens_per_second_per_gpu:6.10059928894043 peak_memory_active:27.452006340026855 peak_memory_alloc:27.452006340026855 peak_memory_reserved:39.03515625
Step 5 | loss:1.4565016031265259 lr:2e-05 tokens_per_second_per_gpu:10.525087356567383 peak_memory_active:27.42187261581421 peak_memory_alloc:27.42187261581421 peak_memory_reserved:39.03515625
Step 6 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.103992462158203 peak_memory_active:27.448771476745605 peak_memory_alloc:27.448771476745605 peak_memory_reserved:39.03515625
Step 7 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.136734962463379 peak_memory_active:27.407429218292236 peak_memory_alloc:27.407429218292236 peak_memory_reserved:39.03515625
Step 8 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.029473304748535 peak_memory_active:27.543721675872803 peak_memory_alloc:27.543721675872803 peak_memory_reserved:39.03515625
Step 9 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.074372291564941 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625
Step 10 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.008955955505371 peak_memory_active:27.40485429763794 peak_memory_alloc:27.40485429763794 peak_memory_reserved:39.03515625
Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.958040237426758 peak_memory_active:27.411057949066162 peak_memory_alloc:27.411057949066162 peak_memory_reserved:39.03515625
Step 12 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.930673599243164 peak_memory_active:27.410027980804443 peak_memory_alloc:27.410027980804443 peak_memory_reserved:39.03515625
Step 13 | loss:nan lr:2e-05 tokens_per_second_per_gpu:13.09176254272461 peak_memory_active:27.406914234161377 peak_memory_alloc:27.406914234161377 peak_memory_reserved:39.03515625
Step 14 | loss:nan lr:2e-05 tokens_per_second_per_gpu:1.858479619026184 peak_memory_active:27.490173816680908 peak_memory_alloc:27.490173816680908 peak_memory_reserved:39.03515625
Step 15 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.682432651519775 peak_memory_active:27.47649908065796 peak_memory_alloc:27.47649908065796 peak_memory_reserved:39.03515625
Step 16 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.597535133361816 peak_memory_active:27.44224739074707 peak_memory_alloc:27.44224739074707 peak_memory_reserved:39.03515625
Step 17 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.915543079376221 peak_memory_active:27.41466236114502 peak_memory_alloc:27.41466236114502 peak_memory_reserved:39.03515625
Step 18 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.867782592773438 peak_memory_active:27.400733947753906 peak_memory_alloc:27.400733947753906 peak_memory_reserved:39.03515625
Step 19 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.128471374511719 peak_memory_active:27.4446063041687 peak_memory_alloc:27.4446063041687 peak_memory_reserved:39.03515625
Step 20 | loss:nan lr:2e-05 tokens_per_second_per_gpu:14.666069030761719 peak_memory_active:27.431589126586914 peak_memory_alloc:27.431589126586914 peak_memory_reserved:39.03515625
Step 21 | loss:nan lr:2e-05 tokens_per_second_per_gpu:10.821459770202637 peak_memory_active:27.424962043762207 peak_memory_alloc:27.424962043762207 peak_memory_reserved:39.037109375
Step 22 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.869871139526367 peak_memory_active:27.419812202453613 peak_memory_alloc:27.419812202453613 peak_memory_reserved:39.037109375


└─ $ tune_logs_tps log_1747916676.txt
tps_avg: 8.70872  n= 24
peak_memory_alloc max   27.54
peak_memory_reser max   39.04

[ghstack-poisoned]

pytorch-bot · 2025-05-21T14:59:35Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2755

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8864b5d with merge base 8195b2e ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Enablement copy pasted from torchtitan. There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty) ``` out = _grouped_mm(h, w2, offs=offsets) out[offsets[-1] :].zero_() ``` scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index. To prevent copying unitialized data - zero_() after grouped_mm for now. [ghstack-poisoned]

pbontrager · 2025-05-21T20:41:52Z

torchtune/models/llama4/_component_builders.py

@@ -180,6 +181,7 @@ def llama4_decoder(
    num_experts: int = 16,
    experts_per_token: int = 1,
    use_shared_expert: bool = True,
+    use_grouped_mm: bool = True,


Do we actually need to expose this as an option? Are there cases where it's supported and not the fastest solution?

I exposed it only for debugging. grouped_mm way was added recently and there could be some issues (NaNs etc.) . I think it will be good to have a switch to test the default version. Once everything is tested - to remove it. Alternatively we can have a global config for this, and not to change the Modules interface.

Yeah I am OK to keep it configurable for now

Most of the NaN issues are resolved, right? Ofc due to the padding I suspect there will be slightly different numerics, but would still be good to explicitly confirm that we get similar loss curves out of the grouped_mm vs non-grouped_mm implementations. Personally I would lean towards just taking this out (especially since it's hardcoded in the top-level Scout builder anyways) cause we already have enough args as it is. But if you wanna keep it in until things stabilize I don't mind too much

Enablement copy pasted from torchtitan. There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty) ``` out = _grouped_mm(h, w2, offs=offsets) out[offsets[-1] :].zero_() ``` scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index. To prevent copying unitialized data - zero_() after grouped_mm for now. [ghstack-poisoned]

codecov-commenter · 2025-05-22T13:09:54Z

Codecov Report

Attention: Patch coverage is 18.34862% with 89 lines in your changes missing coverage. Please review.

Please upload report for BASE (gh/IvanKobzarev/5/base@428f301). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
torchtune/modules/moe/indices.py	0.00%	57 Missing ⚠️
torchtune/modules/moe/moe.py	5.26%	18 Missing ⚠️
torchtune/modules/moe/experts.py	44.44%	10 Missing ⚠️
recipes/full_finetune_distributed.py	0.00%	3 Missing ⚠️
torchtune/modules/moe/config.py	90.00%	1 Missing ⚠️

Additional details and impacted files

@@                    Coverage Diff                    @@
##             gh/IvanKobzarev/5/base    #2755   +/-   ##
=========================================================
  Coverage                          ?   59.91%           
=========================================================
  Files                             ?      437           
  Lines                             ?    26848           
  Branches                          ?        0           
=========================================================
  Hits                              ?    16085           
  Misses                            ?    10763           
  Partials                          ?        0

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Enablement copy pasted from torchtitan. There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty) ``` out = _grouped_mm(h, w2, offs=offsets) out[offsets[-1] :].zero_() ``` scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index. To prevent copying unitialized data - zero_() after grouped_mm for now. grouped_mm in eager: eager tps_avg: 3.2182 -> 4.44 grouped_mm in compile: compile tps_avg: 4.33463 -> 8.70872 Observing NaN loss in compile :( with/without grouped_mm - will be debugging this separately ``` no_compile + no grouped_mm: BASELINE: Step 1 | loss:4.149928092956543 lr:2e-05 tokens_per_second_per_gpu:2.203409433364868 peak_memory_active:28.82542610168457 peak_memory_alloc:28.82542610168457 peak_memory_reserved:40.197265625 Step 2 | loss:2.3132553100585938 lr:2e-05 tokens_per_second_per_gpu:2.492661952972412 peak_memory_active:28.804787158966064 peak_memory_alloc:28.804787158966064 peak_memory_reserved:40.251953125 Step 3 | loss:1.6078141927719116 lr:2e-05 tokens_per_second_per_gpu:3.8224945068359375 peak_memory_active:28.853728771209717 peak_memory_alloc:28.853728771209717 peak_memory_reserved:40.251953125 Step 4 | loss:1.4519065618515015 lr:2e-05 tokens_per_second_per_gpu:3.001920700073242 peak_memory_active:28.84491729736328 peak_memory_alloc:28.84491729736328 peak_memory_reserved:40.251953125 Step 5 | loss:1.2131776809692383 lr:2e-05 tokens_per_second_per_gpu:3.3134357929229736 peak_memory_active:28.804935932159424 peak_memory_alloc:28.804935932159424 peak_memory_reserved:40.251953125 Step 6 | loss:1.411360263824463 lr:2e-05 tokens_per_second_per_gpu:4.279047966003418 peak_memory_active:28.83890438079834 peak_memory_alloc:28.83890438079834 peak_memory_reserved:40.251953125 Step 7 | loss:1.1743241548538208 lr:2e-05 tokens_per_second_per_gpu:3.129912853240967 peak_memory_active:28.785582065582275 peak_memory_alloc:28.785582065582275 peak_memory_reserved:40.251953125 Step 8 | loss:1.3950272798538208 lr:2e-05 tokens_per_second_per_gpu:3.4216036796569824 peak_memory_active:28.967872619628906 peak_memory_alloc:28.967872619628906 peak_memory_reserved:40.251953125 Step 9 | loss:1.2500101327896118 lr:2e-05 tokens_per_second_per_gpu:3.410867214202881 peak_memory_active:28.854501247406006 peak_memory_alloc:28.854501247406006 peak_memory_reserved:40.251953125 Step 10 | loss:1.1264036893844604 lr:2e-05 tokens_per_second_per_gpu:3.1065995693206787 peak_memory_active:28.782616138458252 peak_memory_alloc:28.782616138458252 peak_memory_reserved:40.251953125 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747910831.txt tps_avg: 3.2182 n= 10 peak_memory_alloc max 28.97 peak_memory_reser max 40.25 no_compile + grouped_mm: ______ EAGER_GROUPED_MM: └─ $ cat log_1747915295.txt Step 1 | loss:4.163714408874512 lr:2e-05 tokens_per_second_per_gpu:2.171847343444824 peak_memory_active:27.444169998168945 peak_memory_alloc:27.444169998168945 peak_memory_reserved:38.9609375 Step 2 | loss:2.283238172531128 lr:2e-05 tokens_per_second_per_gpu:3.158111810684204 peak_memory_active:27.428916931152344 peak_memory_alloc:27.428916931152344 peak_memory_reserved:39.017578125 Step 3 | loss:1.5704680681228638 lr:2e-05 tokens_per_second_per_gpu:4.869690418243408 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 4 | loss:1.4485288858413696 lr:2e-05 tokens_per_second_per_gpu:4.026949405670166 peak_memory_active:27.46179723739624 peak_memory_alloc:27.46179723739624 peak_memory_reserved:39.017578125 Step 5 | loss:1.1918020248413086 lr:2e-05 tokens_per_second_per_gpu:4.502264976501465 peak_memory_active:27.428401947021484 peak_memory_alloc:27.428401947021484 peak_memory_reserved:39.017578125 Step 6 | loss:1.3454418182373047 lr:2e-05 tokens_per_second_per_gpu:5.665891170501709 peak_memory_active:27.456215858459473 peak_memory_alloc:27.456215858459473 peak_memory_reserved:39.017578125 Step 7 | loss:1.185951590538025 lr:2e-05 tokens_per_second_per_gpu:4.356126308441162 peak_memory_active:27.412775993347168 peak_memory_alloc:27.412775993347168 peak_memory_reserved:39.017578125 Step 8 | loss:1.3995944261550903 lr:2e-05 tokens_per_second_per_gpu:4.84677267074585 peak_memory_active:27.562548637390137 peak_memory_alloc:27.562548637390137 peak_memory_reserved:39.017578125 Step 9 | loss:1.2516542673110962 lr:2e-05 tokens_per_second_per_gpu:4.785239219665527 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 10 | loss:1.1245404481887817 lr:2e-05 tokens_per_second_per_gpu:4.080092906951904 peak_memory_active:27.41029644012451 peak_memory_alloc:27.41029644012451 peak_memory_reserved:39.017578125 Step 11 | loss:1.2505868673324585 lr:2e-05 tokens_per_second_per_gpu:4.96058988571167 peak_memory_active:27.416271209716797 peak_memory_alloc:27.416271209716797 peak_memory_reserved:39.0234375 Step 12 | loss:1.0976852178573608 lr:2e-05 tokens_per_second_per_gpu:4.555960655212402 peak_memory_active:27.415279388427734 peak_memory_alloc:27.415279388427734 peak_memory_reserved:39.0234375 Step 13 | loss:1.176978349685669 lr:2e-05 tokens_per_second_per_gpu:5.802102565765381 peak_memory_active:27.413253784179688 peak_memory_alloc:27.413253784179688 peak_memory_reserved:39.0234375 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915295.txt tps_avg: 4.44474 n= 13 peak_memory_alloc max 27.56 peak_memory_reser max 39.02 compile + no_grouped_mm: COMPILED BASELINE: └─ $ cat log_1747915995.txt Step 1 | loss:4.1595354080200195 lr:2e-05 tokens_per_second_per_gpu:0.42956066131591797 peak_memory_active:28.81977891921997 peak_memory_alloc:28.81977891921997 peak_memory_reserved:40.236328125 Step 2 | loss:2.2778103351593018 lr:2e-05 tokens_per_second_per_gpu:2.819465160369873 peak_memory_active:28.79981756210327 peak_memory_alloc:28.79981756210327 peak_memory_reserved:40.2890625 Step 3 | loss:1.786258578300476 lr:2e-05 tokens_per_second_per_gpu:5.622256755828857 peak_memory_active:28.846065998077393 peak_memory_alloc:28.846065998077393 peak_memory_reserved:40.2890625 Step 4 | loss:1.600595474243164 lr:2e-05 tokens_per_second_per_gpu:2.996225357055664 peak_memory_active:28.837817668914795 peak_memory_alloc:28.837817668914795 peak_memory_reserved:40.2890625 Step 5 | loss:1.627750039100647 lr:2e-05 tokens_per_second_per_gpu:6.902056694030762 peak_memory_active:28.79964590072632 peak_memory_alloc:28.79964590072632 peak_memory_reserved:40.2890625 Step 6 | loss:1.6751512289047241 lr:2e-05 tokens_per_second_per_gpu:3.8652405738830566 peak_memory_active:28.833858966827393 peak_memory_alloc:28.833858966827393 peak_memory_reserved:40.2890625 Step 7 | loss:1.5953153371810913 lr:2e-05 tokens_per_second_per_gpu:2.8824896812438965 peak_memory_active:28.781311511993408 peak_memory_alloc:28.781311511993408 peak_memory_reserved:40.2890625 Step 8 | loss:1.6390446424484253 lr:2e-05 tokens_per_second_per_gpu:4.305164813995361 peak_memory_active:28.948315620422363 peak_memory_alloc:28.948315620422363 peak_memory_reserved:40.2890625 Step 9 | loss:1.53915536403656 lr:2e-05 tokens_per_second_per_gpu:4.096757888793945 peak_memory_active:28.84673547744751 peak_memory_alloc:28.84673547744751 peak_memory_reserved:40.2890625 Step 10 | loss:1.4715595245361328 lr:2e-05 tokens_per_second_per_gpu:6.167354106903076 peak_memory_active:28.778648853302002 peak_memory_alloc:28.778648853302002 peak_memory_reserved:40.2890625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.351371765136719 peak_memory_active:28.787389278411865 peak_memory_alloc:28.787389278411865 peak_memory_reserved:40.2890625 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915995.txt tps_avg: 4.33463 n= 12 peak_memory_alloc max 28.95 peak_memory_reser max 40.29 COMPILED + GROUPED_MM └─ $ cat log_1747916676.txt Step 1 | loss:4.163997650146484 lr:2e-05 tokens_per_second_per_gpu:0.9783019423484802 peak_memory_active:27.437111377716064 peak_memory_alloc:27.437111377716064 peak_memory_reserved:38.978515625 Step 2 | loss:2.2560360431671143 lr:2e-05 tokens_per_second_per_gpu:2.861786365509033 peak_memory_active:27.42238759994507 peak_memory_alloc:27.42238759994507 peak_memory_reserved:39.03515625 Step 3 | loss:1.8066692352294922 lr:2e-05 tokens_per_second_per_gpu:7.731247901916504 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 4 | loss:1.7610063552856445 lr:2e-05 tokens_per_second_per_gpu:6.10059928894043 peak_memory_active:27.452006340026855 peak_memory_alloc:27.452006340026855 peak_memory_reserved:39.03515625 Step 5 | loss:1.4565016031265259 lr:2e-05 tokens_per_second_per_gpu:10.525087356567383 peak_memory_active:27.42187261581421 peak_memory_alloc:27.42187261581421 peak_memory_reserved:39.03515625 Step 6 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.103992462158203 peak_memory_active:27.448771476745605 peak_memory_alloc:27.448771476745605 peak_memory_reserved:39.03515625 Step 7 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.136734962463379 peak_memory_active:27.407429218292236 peak_memory_alloc:27.407429218292236 peak_memory_reserved:39.03515625 Step 8 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.029473304748535 peak_memory_active:27.543721675872803 peak_memory_alloc:27.543721675872803 peak_memory_reserved:39.03515625 Step 9 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.074372291564941 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 10 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.008955955505371 peak_memory_active:27.40485429763794 peak_memory_alloc:27.40485429763794 peak_memory_reserved:39.03515625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.958040237426758 peak_memory_active:27.411057949066162 peak_memory_alloc:27.411057949066162 peak_memory_reserved:39.03515625 Step 12 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.930673599243164 peak_memory_active:27.410027980804443 peak_memory_alloc:27.410027980804443 peak_memory_reserved:39.03515625 Step 13 | loss:nan lr:2e-05 tokens_per_second_per_gpu:13.09176254272461 peak_memory_active:27.406914234161377 peak_memory_alloc:27.406914234161377 peak_memory_reserved:39.03515625 Step 14 | loss:nan lr:2e-05 tokens_per_second_per_gpu:1.858479619026184 peak_memory_active:27.490173816680908 peak_memory_alloc:27.490173816680908 peak_memory_reserved:39.03515625 Step 15 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.682432651519775 peak_memory_active:27.47649908065796 peak_memory_alloc:27.47649908065796 peak_memory_reserved:39.03515625 Step 16 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.597535133361816 peak_memory_active:27.44224739074707 peak_memory_alloc:27.44224739074707 peak_memory_reserved:39.03515625 Step 17 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.915543079376221 peak_memory_active:27.41466236114502 peak_memory_alloc:27.41466236114502 peak_memory_reserved:39.03515625 Step 18 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.867782592773438 peak_memory_active:27.400733947753906 peak_memory_alloc:27.400733947753906 peak_memory_reserved:39.03515625 Step 19 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.128471374511719 peak_memory_active:27.4446063041687 peak_memory_alloc:27.4446063041687 peak_memory_reserved:39.03515625 Step 20 | loss:nan lr:2e-05 tokens_per_second_per_gpu:14.666069030761719 peak_memory_active:27.431589126586914 peak_memory_alloc:27.431589126586914 peak_memory_reserved:39.03515625 Step 21 | loss:nan lr:2e-05 tokens_per_second_per_gpu:10.821459770202637 peak_memory_active:27.424962043762207 peak_memory_alloc:27.424962043762207 peak_memory_reserved:39.037109375 Step 22 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.869871139526367 peak_memory_active:27.419812202453613 peak_memory_alloc:27.419812202453613 peak_memory_reserved:39.037109375 └─ $ tune_logs_tps log_1747916676.txt tps_avg: 8.70872 n= 24 peak_memory_alloc max 27.54 peak_memory_reser max 39.04 ``` [ghstack-poisoned]

ghstack-source-id: 88f77b2 Pull Request resolved: #2755

ebsmothers

Looks good! Left a handful of comments, lmk if anything is unclear. Excited to get grouped_mm landed!

ebsmothers · 2025-05-29T21:54:47Z

torchtune/modules/moe/indices.py

+    return permuted_indices, m_sizes, m_offsets.to(torch.int32)
+
+
+# Below is for testing only


Can we put the below code in a separate unit test? E.g. tests/torchtune/modules/moe/test_indices.py?

ebsmothers · 2025-05-29T21:55:14Z

torchtune/modules/moe/indices.py

+    print(f"tokens_per_expert_group = {tokens_per_expert_group}")
+    print(f"total_tokens_per_expert = {total_tokens_per_expert}")
+    print(f"m_sizes = {m_sizes}")
+    print(f"m_offsets = {m_offsets}")
+    print(f"permuted_indices = {permuted_indices_gpu[:sum(m_sizes).item()]}")


Can remove these too and just rely on asserts once we move to a test

ebsmothers · 2025-05-29T22:05:37Z

torchtune/modules/moe/indices.py

+import triton.language as tl
+
+
+__all__ = ["generate_permute_indices"]


nit: we usually put __all__ in the __init__.py file so if you want to expose as a public API you should do it there (but personally I also don't mind keeping this as a private API)

Yes, I think better to keep it private.

ebsmothers · 2025-05-29T22:07:46Z

torchtune/models/llama4/_component_builders.py

@@ -180,6 +181,7 @@ def llama4_decoder(
    num_experts: int = 16,
    experts_per_token: int = 1,
    use_shared_expert: bool = True,
+    use_grouped_mm: bool = True,


Yeah I am OK to keep it configurable for now

ebsmothers · 2025-05-29T22:11:21Z

torchtune/modules/moe/indices.py

+
+
+# reference
+def fill_indices_cpu(


Is this just for testing purposes? Or do we actually need it for e.g. CPU training? (Though in that case we probably can't use flex for Llama4 anyways)

I think it is more for a reference and testing. It is not used in llama4 recipes. But I think better to keep it as it is easier to read and understand what kernel does.

ebsmothers · 2025-05-29T23:05:17Z

torchtune/modules/moe/experts.py

+        if num_tokens_per_expert is not None:
+            # https://github.com/pytorch/pytorch/pull/150374
+            # NOTE: torch._gouped_mm requires bf16 dtypes
+            #       and shapes to be multiple of 8


nit: elsewhere it says multiple of 16

ebsmothers · 2025-05-29T23:06:48Z

torchtune/modules/moe/experts.py

+            offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32)
+            # grouped mm between a 2D tensor and a 3D tensor
+            assert x.dim() == 2
+        else:


So this case is for expert choice?

ebsmothers · 2025-05-29T23:11:11Z

torchtune/modules/moe/moe.py

+                    self.experts.num_experts,
+                    1,


This is just because we are fully replicating the experts over all ranks, right?

ebsmothers · 2025-05-29T23:15:38Z

torchtune/modules/moe/indices.py

+        From: |       rank 0      |       rank 1      |
+        To:   | E0 | E1 | E2 | E3 | E0 | E1 | E2 | E3 |


Noob q on this diagram: does this mean that we actually have the same expert on multiple ranks (i.e. there are only 4 experts total)? I ask because the arg experts_per_rank seems to imply that we actually have different experts on different ranks (I know it's moot for the implementation in its current form anyways so not a huge deal).

In our case each rank runs all the experts, but all mms happen in parallel.
So experts_per_rank is always 16 (all experts).

So we are not bucketing experts by ranks. I think this logic will be similar, to what TorchRec does with embeddings, basically each expert works as a lookup Table.

ebsmothers · 2025-05-29T23:16:26Z

torchtune/modules/moe/indices.py

@@ -0,0 +1,353 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.


We should also add an acknowledgment to the original torchtitan code somewhere. @lessw2020 lmk if any preference on the best way to do this

ebsmothers

Looks good! Left a handful of comments, lmk if anything is unclear. Excited to get grouped_mm landed!

Enablement copy pasted from torchtitan. There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty) ``` out = _grouped_mm(h, w2, offs=offsets) out[offsets[-1] :].zero_() ``` scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index. To prevent copying unitialized data - zero_() after grouped_mm for now. grouped_mm in eager: eager tps_avg: 3.2182 -> 4.44 grouped_mm in compile: compile tps_avg: 4.33463 -> 8.70872 Observing NaN loss in compile :( with/without grouped_mm - will be debugging this separately ``` no_compile + no grouped_mm: BASELINE: Step 1 | loss:4.149928092956543 lr:2e-05 tokens_per_second_per_gpu:2.203409433364868 peak_memory_active:28.82542610168457 peak_memory_alloc:28.82542610168457 peak_memory_reserved:40.197265625 Step 2 | loss:2.3132553100585938 lr:2e-05 tokens_per_second_per_gpu:2.492661952972412 peak_memory_active:28.804787158966064 peak_memory_alloc:28.804787158966064 peak_memory_reserved:40.251953125 Step 3 | loss:1.6078141927719116 lr:2e-05 tokens_per_second_per_gpu:3.8224945068359375 peak_memory_active:28.853728771209717 peak_memory_alloc:28.853728771209717 peak_memory_reserved:40.251953125 Step 4 | loss:1.4519065618515015 lr:2e-05 tokens_per_second_per_gpu:3.001920700073242 peak_memory_active:28.84491729736328 peak_memory_alloc:28.84491729736328 peak_memory_reserved:40.251953125 Step 5 | loss:1.2131776809692383 lr:2e-05 tokens_per_second_per_gpu:3.3134357929229736 peak_memory_active:28.804935932159424 peak_memory_alloc:28.804935932159424 peak_memory_reserved:40.251953125 Step 6 | loss:1.411360263824463 lr:2e-05 tokens_per_second_per_gpu:4.279047966003418 peak_memory_active:28.83890438079834 peak_memory_alloc:28.83890438079834 peak_memory_reserved:40.251953125 Step 7 | loss:1.1743241548538208 lr:2e-05 tokens_per_second_per_gpu:3.129912853240967 peak_memory_active:28.785582065582275 peak_memory_alloc:28.785582065582275 peak_memory_reserved:40.251953125 Step 8 | loss:1.3950272798538208 lr:2e-05 tokens_per_second_per_gpu:3.4216036796569824 peak_memory_active:28.967872619628906 peak_memory_alloc:28.967872619628906 peak_memory_reserved:40.251953125 Step 9 | loss:1.2500101327896118 lr:2e-05 tokens_per_second_per_gpu:3.410867214202881 peak_memory_active:28.854501247406006 peak_memory_alloc:28.854501247406006 peak_memory_reserved:40.251953125 Step 10 | loss:1.1264036893844604 lr:2e-05 tokens_per_second_per_gpu:3.1065995693206787 peak_memory_active:28.782616138458252 peak_memory_alloc:28.782616138458252 peak_memory_reserved:40.251953125 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747910831.txt tps_avg: 3.2182 n= 10 peak_memory_alloc max 28.97 peak_memory_reser max 40.25 no_compile + grouped_mm: ______ EAGER_GROUPED_MM: └─ $ cat log_1747915295.txt Step 1 | loss:4.163714408874512 lr:2e-05 tokens_per_second_per_gpu:2.171847343444824 peak_memory_active:27.444169998168945 peak_memory_alloc:27.444169998168945 peak_memory_reserved:38.9609375 Step 2 | loss:2.283238172531128 lr:2e-05 tokens_per_second_per_gpu:3.158111810684204 peak_memory_active:27.428916931152344 peak_memory_alloc:27.428916931152344 peak_memory_reserved:39.017578125 Step 3 | loss:1.5704680681228638 lr:2e-05 tokens_per_second_per_gpu:4.869690418243408 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 4 | loss:1.4485288858413696 lr:2e-05 tokens_per_second_per_gpu:4.026949405670166 peak_memory_active:27.46179723739624 peak_memory_alloc:27.46179723739624 peak_memory_reserved:39.017578125 Step 5 | loss:1.1918020248413086 lr:2e-05 tokens_per_second_per_gpu:4.502264976501465 peak_memory_active:27.428401947021484 peak_memory_alloc:27.428401947021484 peak_memory_reserved:39.017578125 Step 6 | loss:1.3454418182373047 lr:2e-05 tokens_per_second_per_gpu:5.665891170501709 peak_memory_active:27.456215858459473 peak_memory_alloc:27.456215858459473 peak_memory_reserved:39.017578125 Step 7 | loss:1.185951590538025 lr:2e-05 tokens_per_second_per_gpu:4.356126308441162 peak_memory_active:27.412775993347168 peak_memory_alloc:27.412775993347168 peak_memory_reserved:39.017578125 Step 8 | loss:1.3995944261550903 lr:2e-05 tokens_per_second_per_gpu:4.84677267074585 peak_memory_active:27.562548637390137 peak_memory_alloc:27.562548637390137 peak_memory_reserved:39.017578125 Step 9 | loss:1.2516542673110962 lr:2e-05 tokens_per_second_per_gpu:4.785239219665527 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 10 | loss:1.1245404481887817 lr:2e-05 tokens_per_second_per_gpu:4.080092906951904 peak_memory_active:27.41029644012451 peak_memory_alloc:27.41029644012451 peak_memory_reserved:39.017578125 Step 11 | loss:1.2505868673324585 lr:2e-05 tokens_per_second_per_gpu:4.96058988571167 peak_memory_active:27.416271209716797 peak_memory_alloc:27.416271209716797 peak_memory_reserved:39.0234375 Step 12 | loss:1.0976852178573608 lr:2e-05 tokens_per_second_per_gpu:4.555960655212402 peak_memory_active:27.415279388427734 peak_memory_alloc:27.415279388427734 peak_memory_reserved:39.0234375 Step 13 | loss:1.176978349685669 lr:2e-05 tokens_per_second_per_gpu:5.802102565765381 peak_memory_active:27.413253784179688 peak_memory_alloc:27.413253784179688 peak_memory_reserved:39.0234375 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915295.txt tps_avg: 4.44474 n= 13 peak_memory_alloc max 27.56 peak_memory_reser max 39.02 compile + no_grouped_mm: COMPILED BASELINE: └─ $ cat log_1747915995.txt Step 1 | loss:4.1595354080200195 lr:2e-05 tokens_per_second_per_gpu:0.42956066131591797 peak_memory_active:28.81977891921997 peak_memory_alloc:28.81977891921997 peak_memory_reserved:40.236328125 Step 2 | loss:2.2778103351593018 lr:2e-05 tokens_per_second_per_gpu:2.819465160369873 peak_memory_active:28.79981756210327 peak_memory_alloc:28.79981756210327 peak_memory_reserved:40.2890625 Step 3 | loss:1.786258578300476 lr:2e-05 tokens_per_second_per_gpu:5.622256755828857 peak_memory_active:28.846065998077393 peak_memory_alloc:28.846065998077393 peak_memory_reserved:40.2890625 Step 4 | loss:1.600595474243164 lr:2e-05 tokens_per_second_per_gpu:2.996225357055664 peak_memory_active:28.837817668914795 peak_memory_alloc:28.837817668914795 peak_memory_reserved:40.2890625 Step 5 | loss:1.627750039100647 lr:2e-05 tokens_per_second_per_gpu:6.902056694030762 peak_memory_active:28.79964590072632 peak_memory_alloc:28.79964590072632 peak_memory_reserved:40.2890625 Step 6 | loss:1.6751512289047241 lr:2e-05 tokens_per_second_per_gpu:3.8652405738830566 peak_memory_active:28.833858966827393 peak_memory_alloc:28.833858966827393 peak_memory_reserved:40.2890625 Step 7 | loss:1.5953153371810913 lr:2e-05 tokens_per_second_per_gpu:2.8824896812438965 peak_memory_active:28.781311511993408 peak_memory_alloc:28.781311511993408 peak_memory_reserved:40.2890625 Step 8 | loss:1.6390446424484253 lr:2e-05 tokens_per_second_per_gpu:4.305164813995361 peak_memory_active:28.948315620422363 peak_memory_alloc:28.948315620422363 peak_memory_reserved:40.2890625 Step 9 | loss:1.53915536403656 lr:2e-05 tokens_per_second_per_gpu:4.096757888793945 peak_memory_active:28.84673547744751 peak_memory_alloc:28.84673547744751 peak_memory_reserved:40.2890625 Step 10 | loss:1.4715595245361328 lr:2e-05 tokens_per_second_per_gpu:6.167354106903076 peak_memory_active:28.778648853302002 peak_memory_alloc:28.778648853302002 peak_memory_reserved:40.2890625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.351371765136719 peak_memory_active:28.787389278411865 peak_memory_alloc:28.787389278411865 peak_memory_reserved:40.2890625 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915995.txt tps_avg: 4.33463 n= 12 peak_memory_alloc max 28.95 peak_memory_reser max 40.29 COMPILED + GROUPED_MM └─ $ cat log_1747916676.txt Step 1 | loss:4.163997650146484 lr:2e-05 tokens_per_second_per_gpu:0.9783019423484802 peak_memory_active:27.437111377716064 peak_memory_alloc:27.437111377716064 peak_memory_reserved:38.978515625 Step 2 | loss:2.2560360431671143 lr:2e-05 tokens_per_second_per_gpu:2.861786365509033 peak_memory_active:27.42238759994507 peak_memory_alloc:27.42238759994507 peak_memory_reserved:39.03515625 Step 3 | loss:1.8066692352294922 lr:2e-05 tokens_per_second_per_gpu:7.731247901916504 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 4 | loss:1.7610063552856445 lr:2e-05 tokens_per_second_per_gpu:6.10059928894043 peak_memory_active:27.452006340026855 peak_memory_alloc:27.452006340026855 peak_memory_reserved:39.03515625 Step 5 | loss:1.4565016031265259 lr:2e-05 tokens_per_second_per_gpu:10.525087356567383 peak_memory_active:27.42187261581421 peak_memory_alloc:27.42187261581421 peak_memory_reserved:39.03515625 Step 6 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.103992462158203 peak_memory_active:27.448771476745605 peak_memory_alloc:27.448771476745605 peak_memory_reserved:39.03515625 Step 7 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.136734962463379 peak_memory_active:27.407429218292236 peak_memory_alloc:27.407429218292236 peak_memory_reserved:39.03515625 Step 8 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.029473304748535 peak_memory_active:27.543721675872803 peak_memory_alloc:27.543721675872803 peak_memory_reserved:39.03515625 Step 9 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.074372291564941 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 10 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.008955955505371 peak_memory_active:27.40485429763794 peak_memory_alloc:27.40485429763794 peak_memory_reserved:39.03515625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.958040237426758 peak_memory_active:27.411057949066162 peak_memory_alloc:27.411057949066162 peak_memory_reserved:39.03515625 Step 12 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.930673599243164 peak_memory_active:27.410027980804443 peak_memory_alloc:27.410027980804443 peak_memory_reserved:39.03515625 Step 13 | loss:nan lr:2e-05 tokens_per_second_per_gpu:13.09176254272461 peak_memory_active:27.406914234161377 peak_memory_alloc:27.406914234161377 peak_memory_reserved:39.03515625 Step 14 | loss:nan lr:2e-05 tokens_per_second_per_gpu:1.858479619026184 peak_memory_active:27.490173816680908 peak_memory_alloc:27.490173816680908 peak_memory_reserved:39.03515625 Step 15 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.682432651519775 peak_memory_active:27.47649908065796 peak_memory_alloc:27.47649908065796 peak_memory_reserved:39.03515625 Step 16 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.597535133361816 peak_memory_active:27.44224739074707 peak_memory_alloc:27.44224739074707 peak_memory_reserved:39.03515625 Step 17 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.915543079376221 peak_memory_active:27.41466236114502 peak_memory_alloc:27.41466236114502 peak_memory_reserved:39.03515625 Step 18 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.867782592773438 peak_memory_active:27.400733947753906 peak_memory_alloc:27.400733947753906 peak_memory_reserved:39.03515625 Step 19 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.128471374511719 peak_memory_active:27.4446063041687 peak_memory_alloc:27.4446063041687 peak_memory_reserved:39.03515625 Step 20 | loss:nan lr:2e-05 tokens_per_second_per_gpu:14.666069030761719 peak_memory_active:27.431589126586914 peak_memory_alloc:27.431589126586914 peak_memory_reserved:39.03515625 Step 21 | loss:nan lr:2e-05 tokens_per_second_per_gpu:10.821459770202637 peak_memory_active:27.424962043762207 peak_memory_alloc:27.424962043762207 peak_memory_reserved:39.037109375 Step 22 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.869871139526367 peak_memory_active:27.419812202453613 peak_memory_alloc:27.419812202453613 peak_memory_reserved:39.037109375 └─ $ tune_logs_tps log_1747916676.txt tps_avg: 8.70872 n= 24 peak_memory_alloc max 27.54 peak_memory_reser max 39.04 ``` [ghstack-poisoned]

ghstack-source-id: 8bd646b Pull Request resolved: #2755

Enablement copy pasted from torchtitan. There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty) ``` out = _grouped_mm(h, w2, offs=offsets) out[offsets[-1] :].zero_() ``` scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index. To prevent copying unitialized data - zero_() after grouped_mm for now. grouped_mm in eager: eager tps_avg: 3.2182 -> 4.44 grouped_mm in compile: compile tps_avg: 4.33463 -> 8.70872 Observing NaN loss in compile :( with/without grouped_mm - will be debugging this separately ``` no_compile + no grouped_mm: BASELINE: Step 1 | loss:4.149928092956543 lr:2e-05 tokens_per_second_per_gpu:2.203409433364868 peak_memory_active:28.82542610168457 peak_memory_alloc:28.82542610168457 peak_memory_reserved:40.197265625 Step 2 | loss:2.3132553100585938 lr:2e-05 tokens_per_second_per_gpu:2.492661952972412 peak_memory_active:28.804787158966064 peak_memory_alloc:28.804787158966064 peak_memory_reserved:40.251953125 Step 3 | loss:1.6078141927719116 lr:2e-05 tokens_per_second_per_gpu:3.8224945068359375 peak_memory_active:28.853728771209717 peak_memory_alloc:28.853728771209717 peak_memory_reserved:40.251953125 Step 4 | loss:1.4519065618515015 lr:2e-05 tokens_per_second_per_gpu:3.001920700073242 peak_memory_active:28.84491729736328 peak_memory_alloc:28.84491729736328 peak_memory_reserved:40.251953125 Step 5 | loss:1.2131776809692383 lr:2e-05 tokens_per_second_per_gpu:3.3134357929229736 peak_memory_active:28.804935932159424 peak_memory_alloc:28.804935932159424 peak_memory_reserved:40.251953125 Step 6 | loss:1.411360263824463 lr:2e-05 tokens_per_second_per_gpu:4.279047966003418 peak_memory_active:28.83890438079834 peak_memory_alloc:28.83890438079834 peak_memory_reserved:40.251953125 Step 7 | loss:1.1743241548538208 lr:2e-05 tokens_per_second_per_gpu:3.129912853240967 peak_memory_active:28.785582065582275 peak_memory_alloc:28.785582065582275 peak_memory_reserved:40.251953125 Step 8 | loss:1.3950272798538208 lr:2e-05 tokens_per_second_per_gpu:3.4216036796569824 peak_memory_active:28.967872619628906 peak_memory_alloc:28.967872619628906 peak_memory_reserved:40.251953125 Step 9 | loss:1.2500101327896118 lr:2e-05 tokens_per_second_per_gpu:3.410867214202881 peak_memory_active:28.854501247406006 peak_memory_alloc:28.854501247406006 peak_memory_reserved:40.251953125 Step 10 | loss:1.1264036893844604 lr:2e-05 tokens_per_second_per_gpu:3.1065995693206787 peak_memory_active:28.782616138458252 peak_memory_alloc:28.782616138458252 peak_memory_reserved:40.251953125 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747910831.txt tps_avg: 3.2182 n= 10 peak_memory_alloc max 28.97 peak_memory_reser max 40.25 no_compile + grouped_mm: ______ EAGER_GROUPED_MM: └─ $ cat log_1747915295.txt Step 1 | loss:4.163714408874512 lr:2e-05 tokens_per_second_per_gpu:2.171847343444824 peak_memory_active:27.444169998168945 peak_memory_alloc:27.444169998168945 peak_memory_reserved:38.9609375 Step 2 | loss:2.283238172531128 lr:2e-05 tokens_per_second_per_gpu:3.158111810684204 peak_memory_active:27.428916931152344 peak_memory_alloc:27.428916931152344 peak_memory_reserved:39.017578125 Step 3 | loss:1.5704680681228638 lr:2e-05 tokens_per_second_per_gpu:4.869690418243408 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 4 | loss:1.4485288858413696 lr:2e-05 tokens_per_second_per_gpu:4.026949405670166 peak_memory_active:27.46179723739624 peak_memory_alloc:27.46179723739624 peak_memory_reserved:39.017578125 Step 5 | loss:1.1918020248413086 lr:2e-05 tokens_per_second_per_gpu:4.502264976501465 peak_memory_active:27.428401947021484 peak_memory_alloc:27.428401947021484 peak_memory_reserved:39.017578125 Step 6 | loss:1.3454418182373047 lr:2e-05 tokens_per_second_per_gpu:5.665891170501709 peak_memory_active:27.456215858459473 peak_memory_alloc:27.456215858459473 peak_memory_reserved:39.017578125 Step 7 | loss:1.185951590538025 lr:2e-05 tokens_per_second_per_gpu:4.356126308441162 peak_memory_active:27.412775993347168 peak_memory_alloc:27.412775993347168 peak_memory_reserved:39.017578125 Step 8 | loss:1.3995944261550903 lr:2e-05 tokens_per_second_per_gpu:4.84677267074585 peak_memory_active:27.562548637390137 peak_memory_alloc:27.562548637390137 peak_memory_reserved:39.017578125 Step 9 | loss:1.2516542673110962 lr:2e-05 tokens_per_second_per_gpu:4.785239219665527 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 10 | loss:1.1245404481887817 lr:2e-05 tokens_per_second_per_gpu:4.080092906951904 peak_memory_active:27.41029644012451 peak_memory_alloc:27.41029644012451 peak_memory_reserved:39.017578125 Step 11 | loss:1.2505868673324585 lr:2e-05 tokens_per_second_per_gpu:4.96058988571167 peak_memory_active:27.416271209716797 peak_memory_alloc:27.416271209716797 peak_memory_reserved:39.0234375 Step 12 | loss:1.0976852178573608 lr:2e-05 tokens_per_second_per_gpu:4.555960655212402 peak_memory_active:27.415279388427734 peak_memory_alloc:27.415279388427734 peak_memory_reserved:39.0234375 Step 13 | loss:1.176978349685669 lr:2e-05 tokens_per_second_per_gpu:5.802102565765381 peak_memory_active:27.413253784179688 peak_memory_alloc:27.413253784179688 peak_memory_reserved:39.0234375 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915295.txt tps_avg: 4.44474 n= 13 peak_memory_alloc max 27.56 peak_memory_reser max 39.02 compile + no_grouped_mm: COMPILED BASELINE: └─ $ cat log_1747915995.txt Step 1 | loss:4.1595354080200195 lr:2e-05 tokens_per_second_per_gpu:0.42956066131591797 peak_memory_active:28.81977891921997 peak_memory_alloc:28.81977891921997 peak_memory_reserved:40.236328125 Step 2 | loss:2.2778103351593018 lr:2e-05 tokens_per_second_per_gpu:2.819465160369873 peak_memory_active:28.79981756210327 peak_memory_alloc:28.79981756210327 peak_memory_reserved:40.2890625 Step 3 | loss:1.786258578300476 lr:2e-05 tokens_per_second_per_gpu:5.622256755828857 peak_memory_active:28.846065998077393 peak_memory_alloc:28.846065998077393 peak_memory_reserved:40.2890625 Step 4 | loss:1.600595474243164 lr:2e-05 tokens_per_second_per_gpu:2.996225357055664 peak_memory_active:28.837817668914795 peak_memory_alloc:28.837817668914795 peak_memory_reserved:40.2890625 Step 5 | loss:1.627750039100647 lr:2e-05 tokens_per_second_per_gpu:6.902056694030762 peak_memory_active:28.79964590072632 peak_memory_alloc:28.79964590072632 peak_memory_reserved:40.2890625 Step 6 | loss:1.6751512289047241 lr:2e-05 tokens_per_second_per_gpu:3.8652405738830566 peak_memory_active:28.833858966827393 peak_memory_alloc:28.833858966827393 peak_memory_reserved:40.2890625 Step 7 | loss:1.5953153371810913 lr:2e-05 tokens_per_second_per_gpu:2.8824896812438965 peak_memory_active:28.781311511993408 peak_memory_alloc:28.781311511993408 peak_memory_reserved:40.2890625 Step 8 | loss:1.6390446424484253 lr:2e-05 tokens_per_second_per_gpu:4.305164813995361 peak_memory_active:28.948315620422363 peak_memory_alloc:28.948315620422363 peak_memory_reserved:40.2890625 Step 9 | loss:1.53915536403656 lr:2e-05 tokens_per_second_per_gpu:4.096757888793945 peak_memory_active:28.84673547744751 peak_memory_alloc:28.84673547744751 peak_memory_reserved:40.2890625 Step 10 | loss:1.4715595245361328 lr:2e-05 tokens_per_second_per_gpu:6.167354106903076 peak_memory_active:28.778648853302002 peak_memory_alloc:28.778648853302002 peak_memory_reserved:40.2890625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.351371765136719 peak_memory_active:28.787389278411865 peak_memory_alloc:28.787389278411865 peak_memory_reserved:40.2890625 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915995.txt tps_avg: 4.33463 n= 12 peak_memory_alloc max 28.95 peak_memory_reser max 40.29 COMPILED + GROUPED_MM └─ $ cat log_1747916676.txt Step 1 | loss:4.163997650146484 lr:2e-05 tokens_per_second_per_gpu:0.9783019423484802 peak_memory_active:27.437111377716064 peak_memory_alloc:27.437111377716064 peak_memory_reserved:38.978515625 Step 2 | loss:2.2560360431671143 lr:2e-05 tokens_per_second_per_gpu:2.861786365509033 peak_memory_active:27.42238759994507 peak_memory_alloc:27.42238759994507 peak_memory_reserved:39.03515625 Step 3 | loss:1.8066692352294922 lr:2e-05 tokens_per_second_per_gpu:7.731247901916504 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 4 | loss:1.7610063552856445 lr:2e-05 tokens_per_second_per_gpu:6.10059928894043 peak_memory_active:27.452006340026855 peak_memory_alloc:27.452006340026855 peak_memory_reserved:39.03515625 Step 5 | loss:1.4565016031265259 lr:2e-05 tokens_per_second_per_gpu:10.525087356567383 peak_memory_active:27.42187261581421 peak_memory_alloc:27.42187261581421 peak_memory_reserved:39.03515625 Step 6 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.103992462158203 peak_memory_active:27.448771476745605 peak_memory_alloc:27.448771476745605 peak_memory_reserved:39.03515625 Step 7 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.136734962463379 peak_memory_active:27.407429218292236 peak_memory_alloc:27.407429218292236 peak_memory_reserved:39.03515625 Step 8 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.029473304748535 peak_memory_active:27.543721675872803 peak_memory_alloc:27.543721675872803 peak_memory_reserved:39.03515625 Step 9 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.074372291564941 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 10 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.008955955505371 peak_memory_active:27.40485429763794 peak_memory_alloc:27.40485429763794 peak_memory_reserved:39.03515625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.958040237426758 peak_memory_active:27.411057949066162 peak_memory_alloc:27.411057949066162 peak_memory_reserved:39.03515625 Step 12 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.930673599243164 peak_memory_active:27.410027980804443 peak_memory_alloc:27.410027980804443 peak_memory_reserved:39.03515625 Step 13 | loss:nan lr:2e-05 tokens_per_second_per_gpu:13.09176254272461 peak_memory_active:27.406914234161377 peak_memory_alloc:27.406914234161377 peak_memory_reserved:39.03515625 Step 14 | loss:nan lr:2e-05 tokens_per_second_per_gpu:1.858479619026184 peak_memory_active:27.490173816680908 peak_memory_alloc:27.490173816680908 peak_memory_reserved:39.03515625 Step 15 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.682432651519775 peak_memory_active:27.47649908065796 peak_memory_alloc:27.47649908065796 peak_memory_reserved:39.03515625 Step 16 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.597535133361816 peak_memory_active:27.44224739074707 peak_memory_alloc:27.44224739074707 peak_memory_reserved:39.03515625 Step 17 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.915543079376221 peak_memory_active:27.41466236114502 peak_memory_alloc:27.41466236114502 peak_memory_reserved:39.03515625 Step 18 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.867782592773438 peak_memory_active:27.400733947753906 peak_memory_alloc:27.400733947753906 peak_memory_reserved:39.03515625 Step 19 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.128471374511719 peak_memory_active:27.4446063041687 peak_memory_alloc:27.4446063041687 peak_memory_reserved:39.03515625 Step 20 | loss:nan lr:2e-05 tokens_per_second_per_gpu:14.666069030761719 peak_memory_active:27.431589126586914 peak_memory_alloc:27.431589126586914 peak_memory_reserved:39.03515625 Step 21 | loss:nan lr:2e-05 tokens_per_second_per_gpu:10.821459770202637 peak_memory_active:27.424962043762207 peak_memory_alloc:27.424962043762207 peak_memory_reserved:39.037109375 Step 22 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.869871139526367 peak_memory_active:27.419812202453613 peak_memory_alloc:27.419812202453613 peak_memory_reserved:39.037109375 └─ $ tune_logs_tps log_1747916676.txt tps_avg: 8.70872 n= 24 peak_memory_alloc max 27.54 peak_memory_reser max 39.04 ``` [ghstack-poisoned]

ghstack-source-id: 2045c31 Pull Request resolved: #2755

Enablement copy pasted from torchtitan. There is only one difference - output of grouped_mm is unintialized after offsets[-1] (result of at::empty) ``` out = _grouped_mm(h, w2, offs=offsets) out[offsets[-1] :].zero_() ``` scatter_add does not have ignore indices and default 0 in token_indices will make it copied to the 0 index. To prevent copying unitialized data - zero_() after grouped_mm for now. grouped_mm in eager: eager tps_avg: 3.2182 -> 4.44 grouped_mm in compile: compile tps_avg: 4.33463 -> 8.70872 Observing NaN loss in compile :( with/without grouped_mm - will be debugging this separately ``` no_compile + no grouped_mm: BASELINE: Step 1 | loss:4.149928092956543 lr:2e-05 tokens_per_second_per_gpu:2.203409433364868 peak_memory_active:28.82542610168457 peak_memory_alloc:28.82542610168457 peak_memory_reserved:40.197265625 Step 2 | loss:2.3132553100585938 lr:2e-05 tokens_per_second_per_gpu:2.492661952972412 peak_memory_active:28.804787158966064 peak_memory_alloc:28.804787158966064 peak_memory_reserved:40.251953125 Step 3 | loss:1.6078141927719116 lr:2e-05 tokens_per_second_per_gpu:3.8224945068359375 peak_memory_active:28.853728771209717 peak_memory_alloc:28.853728771209717 peak_memory_reserved:40.251953125 Step 4 | loss:1.4519065618515015 lr:2e-05 tokens_per_second_per_gpu:3.001920700073242 peak_memory_active:28.84491729736328 peak_memory_alloc:28.84491729736328 peak_memory_reserved:40.251953125 Step 5 | loss:1.2131776809692383 lr:2e-05 tokens_per_second_per_gpu:3.3134357929229736 peak_memory_active:28.804935932159424 peak_memory_alloc:28.804935932159424 peak_memory_reserved:40.251953125 Step 6 | loss:1.411360263824463 lr:2e-05 tokens_per_second_per_gpu:4.279047966003418 peak_memory_active:28.83890438079834 peak_memory_alloc:28.83890438079834 peak_memory_reserved:40.251953125 Step 7 | loss:1.1743241548538208 lr:2e-05 tokens_per_second_per_gpu:3.129912853240967 peak_memory_active:28.785582065582275 peak_memory_alloc:28.785582065582275 peak_memory_reserved:40.251953125 Step 8 | loss:1.3950272798538208 lr:2e-05 tokens_per_second_per_gpu:3.4216036796569824 peak_memory_active:28.967872619628906 peak_memory_alloc:28.967872619628906 peak_memory_reserved:40.251953125 Step 9 | loss:1.2500101327896118 lr:2e-05 tokens_per_second_per_gpu:3.410867214202881 peak_memory_active:28.854501247406006 peak_memory_alloc:28.854501247406006 peak_memory_reserved:40.251953125 Step 10 | loss:1.1264036893844604 lr:2e-05 tokens_per_second_per_gpu:3.1065995693206787 peak_memory_active:28.782616138458252 peak_memory_alloc:28.782616138458252 peak_memory_reserved:40.251953125 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747910831.txt tps_avg: 3.2182 n= 10 peak_memory_alloc max 28.97 peak_memory_reser max 40.25 no_compile + grouped_mm: ______ EAGER_GROUPED_MM: └─ $ cat log_1747915295.txt Step 1 | loss:4.163714408874512 lr:2e-05 tokens_per_second_per_gpu:2.171847343444824 peak_memory_active:27.444169998168945 peak_memory_alloc:27.444169998168945 peak_memory_reserved:38.9609375 Step 2 | loss:2.283238172531128 lr:2e-05 tokens_per_second_per_gpu:3.158111810684204 peak_memory_active:27.428916931152344 peak_memory_alloc:27.428916931152344 peak_memory_reserved:39.017578125 Step 3 | loss:1.5704680681228638 lr:2e-05 tokens_per_second_per_gpu:4.869690418243408 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 4 | loss:1.4485288858413696 lr:2e-05 tokens_per_second_per_gpu:4.026949405670166 peak_memory_active:27.46179723739624 peak_memory_alloc:27.46179723739624 peak_memory_reserved:39.017578125 Step 5 | loss:1.1918020248413086 lr:2e-05 tokens_per_second_per_gpu:4.502264976501465 peak_memory_active:27.428401947021484 peak_memory_alloc:27.428401947021484 peak_memory_reserved:39.017578125 Step 6 | loss:1.3454418182373047 lr:2e-05 tokens_per_second_per_gpu:5.665891170501709 peak_memory_active:27.456215858459473 peak_memory_alloc:27.456215858459473 peak_memory_reserved:39.017578125 Step 7 | loss:1.185951590538025 lr:2e-05 tokens_per_second_per_gpu:4.356126308441162 peak_memory_active:27.412775993347168 peak_memory_alloc:27.412775993347168 peak_memory_reserved:39.017578125 Step 8 | loss:1.3995944261550903 lr:2e-05 tokens_per_second_per_gpu:4.84677267074585 peak_memory_active:27.562548637390137 peak_memory_alloc:27.562548637390137 peak_memory_reserved:39.017578125 Step 9 | loss:1.2516542673110962 lr:2e-05 tokens_per_second_per_gpu:4.785239219665527 peak_memory_active:27.469550132751465 peak_memory_alloc:27.469550132751465 peak_memory_reserved:39.017578125 Step 10 | loss:1.1245404481887817 lr:2e-05 tokens_per_second_per_gpu:4.080092906951904 peak_memory_active:27.41029644012451 peak_memory_alloc:27.41029644012451 peak_memory_reserved:39.017578125 Step 11 | loss:1.2505868673324585 lr:2e-05 tokens_per_second_per_gpu:4.96058988571167 peak_memory_active:27.416271209716797 peak_memory_alloc:27.416271209716797 peak_memory_reserved:39.0234375 Step 12 | loss:1.0976852178573608 lr:2e-05 tokens_per_second_per_gpu:4.555960655212402 peak_memory_active:27.415279388427734 peak_memory_alloc:27.415279388427734 peak_memory_reserved:39.0234375 Step 13 | loss:1.176978349685669 lr:2e-05 tokens_per_second_per_gpu:5.802102565765381 peak_memory_active:27.413253784179688 peak_memory_alloc:27.413253784179688 peak_memory_reserved:39.0234375 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915295.txt tps_avg: 4.44474 n= 13 peak_memory_alloc max 27.56 peak_memory_reser max 39.02 compile + no_grouped_mm: COMPILED BASELINE: └─ $ cat log_1747915995.txt Step 1 | loss:4.1595354080200195 lr:2e-05 tokens_per_second_per_gpu:0.42956066131591797 peak_memory_active:28.81977891921997 peak_memory_alloc:28.81977891921997 peak_memory_reserved:40.236328125 Step 2 | loss:2.2778103351593018 lr:2e-05 tokens_per_second_per_gpu:2.819465160369873 peak_memory_active:28.79981756210327 peak_memory_alloc:28.79981756210327 peak_memory_reserved:40.2890625 Step 3 | loss:1.786258578300476 lr:2e-05 tokens_per_second_per_gpu:5.622256755828857 peak_memory_active:28.846065998077393 peak_memory_alloc:28.846065998077393 peak_memory_reserved:40.2890625 Step 4 | loss:1.600595474243164 lr:2e-05 tokens_per_second_per_gpu:2.996225357055664 peak_memory_active:28.837817668914795 peak_memory_alloc:28.837817668914795 peak_memory_reserved:40.2890625 Step 5 | loss:1.627750039100647 lr:2e-05 tokens_per_second_per_gpu:6.902056694030762 peak_memory_active:28.79964590072632 peak_memory_alloc:28.79964590072632 peak_memory_reserved:40.2890625 Step 6 | loss:1.6751512289047241 lr:2e-05 tokens_per_second_per_gpu:3.8652405738830566 peak_memory_active:28.833858966827393 peak_memory_alloc:28.833858966827393 peak_memory_reserved:40.2890625 Step 7 | loss:1.5953153371810913 lr:2e-05 tokens_per_second_per_gpu:2.8824896812438965 peak_memory_active:28.781311511993408 peak_memory_alloc:28.781311511993408 peak_memory_reserved:40.2890625 Step 8 | loss:1.6390446424484253 lr:2e-05 tokens_per_second_per_gpu:4.305164813995361 peak_memory_active:28.948315620422363 peak_memory_alloc:28.948315620422363 peak_memory_reserved:40.2890625 Step 9 | loss:1.53915536403656 lr:2e-05 tokens_per_second_per_gpu:4.096757888793945 peak_memory_active:28.84673547744751 peak_memory_alloc:28.84673547744751 peak_memory_reserved:40.2890625 Step 10 | loss:1.4715595245361328 lr:2e-05 tokens_per_second_per_gpu:6.167354106903076 peak_memory_active:28.778648853302002 peak_memory_alloc:28.778648853302002 peak_memory_reserved:40.2890625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.351371765136719 peak_memory_active:28.787389278411865 peak_memory_alloc:28.787389278411865 peak_memory_reserved:40.2890625 /tmp/torchtune/llama4_17Bx16E/full/logs └─ $ tune_logs_tps log_1747915995.txt tps_avg: 4.33463 n= 12 peak_memory_alloc max 28.95 peak_memory_reser max 40.29 COMPILED + GROUPED_MM └─ $ cat log_1747916676.txt Step 1 | loss:4.163997650146484 lr:2e-05 tokens_per_second_per_gpu:0.9783019423484802 peak_memory_active:27.437111377716064 peak_memory_alloc:27.437111377716064 peak_memory_reserved:38.978515625 Step 2 | loss:2.2560360431671143 lr:2e-05 tokens_per_second_per_gpu:2.861786365509033 peak_memory_active:27.42238759994507 peak_memory_alloc:27.42238759994507 peak_memory_reserved:39.03515625 Step 3 | loss:1.8066692352294922 lr:2e-05 tokens_per_second_per_gpu:7.731247901916504 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 4 | loss:1.7610063552856445 lr:2e-05 tokens_per_second_per_gpu:6.10059928894043 peak_memory_active:27.452006340026855 peak_memory_alloc:27.452006340026855 peak_memory_reserved:39.03515625 Step 5 | loss:1.4565016031265259 lr:2e-05 tokens_per_second_per_gpu:10.525087356567383 peak_memory_active:27.42187261581421 peak_memory_alloc:27.42187261581421 peak_memory_reserved:39.03515625 Step 6 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.103992462158203 peak_memory_active:27.448771476745605 peak_memory_alloc:27.448771476745605 peak_memory_reserved:39.03515625 Step 7 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.136734962463379 peak_memory_active:27.407429218292236 peak_memory_alloc:27.407429218292236 peak_memory_reserved:39.03515625 Step 8 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.029473304748535 peak_memory_active:27.543721675872803 peak_memory_alloc:27.543721675872803 peak_memory_reserved:39.03515625 Step 9 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.074372291564941 peak_memory_active:27.45892906188965 peak_memory_alloc:27.45892906188965 peak_memory_reserved:39.03515625 Step 10 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.008955955505371 peak_memory_active:27.40485429763794 peak_memory_alloc:27.40485429763794 peak_memory_reserved:39.03515625 Step 11 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.958040237426758 peak_memory_active:27.411057949066162 peak_memory_alloc:27.411057949066162 peak_memory_reserved:39.03515625 Step 12 | loss:nan lr:2e-05 tokens_per_second_per_gpu:9.930673599243164 peak_memory_active:27.410027980804443 peak_memory_alloc:27.410027980804443 peak_memory_reserved:39.03515625 Step 13 | loss:nan lr:2e-05 tokens_per_second_per_gpu:13.09176254272461 peak_memory_active:27.406914234161377 peak_memory_alloc:27.406914234161377 peak_memory_reserved:39.03515625 Step 14 | loss:nan lr:2e-05 tokens_per_second_per_gpu:1.858479619026184 peak_memory_active:27.490173816680908 peak_memory_alloc:27.490173816680908 peak_memory_reserved:39.03515625 Step 15 | loss:nan lr:2e-05 tokens_per_second_per_gpu:7.682432651519775 peak_memory_active:27.47649908065796 peak_memory_alloc:27.47649908065796 peak_memory_reserved:39.03515625 Step 16 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.597535133361816 peak_memory_active:27.44224739074707 peak_memory_alloc:27.44224739074707 peak_memory_reserved:39.03515625 Step 17 | loss:nan lr:2e-05 tokens_per_second_per_gpu:4.915543079376221 peak_memory_active:27.41466236114502 peak_memory_alloc:27.41466236114502 peak_memory_reserved:39.03515625 Step 18 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.867782592773438 peak_memory_active:27.400733947753906 peak_memory_alloc:27.400733947753906 peak_memory_reserved:39.03515625 Step 19 | loss:nan lr:2e-05 tokens_per_second_per_gpu:11.128471374511719 peak_memory_active:27.4446063041687 peak_memory_alloc:27.4446063041687 peak_memory_reserved:39.03515625 Step 20 | loss:nan lr:2e-05 tokens_per_second_per_gpu:14.666069030761719 peak_memory_active:27.431589126586914 peak_memory_alloc:27.431589126586914 peak_memory_reserved:39.03515625 Step 21 | loss:nan lr:2e-05 tokens_per_second_per_gpu:10.821459770202637 peak_memory_active:27.424962043762207 peak_memory_alloc:27.424962043762207 peak_memory_reserved:39.037109375 Step 22 | loss:nan lr:2e-05 tokens_per_second_per_gpu:8.869871139526367 peak_memory_active:27.419812202453613 peak_memory_alloc:27.419812202453613 peak_memory_reserved:39.037109375 └─ $ tune_logs_tps log_1747916676.txt tps_avg: 8.70872 n= 24 peak_memory_alloc max 27.54 peak_memory_reser max 39.04 ``` [ghstack-poisoned]

ghstack-source-id: ffd0450 Pull Request resolved: #2755

[llama4] use grouped_mm in moe for sm90

d7b2495

[ghstack-poisoned]

IvanKobzarev mentioned this pull request May 21, 2025

[WIP][DEBUG] llama4 debugging #2756

Open

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 21, 2025

IvanKobzarev requested review from ebsmothers and joecummings May 21, 2025 15:02

pbontrager reviewed May 21, 2025

View reviewed changes

IvanKobzarev added 2 commits May 22, 2025 07:00

IvanKobzarev added a commit that referenced this pull request May 29, 2025

[llama4] use grouped_mm in moe for sm90

f554a2b

ghstack-source-id: 88f77b2 Pull Request resolved: #2755

ebsmothers reviewed May 29, 2025

View reviewed changes

IvanKobzarev added a commit that referenced this pull request Jun 2, 2025

[llama4] use grouped_mm in moe for sm90

1876cab

ghstack-source-id: 8bd646b Pull Request resolved: #2755

IvanKobzarev added a commit that referenced this pull request Jun 2, 2025

[llama4] use grouped_mm in moe for sm90

547ef89

ghstack-source-id: 2045c31 Pull Request resolved: #2755

IvanKobzarev added a commit that referenced this pull request Jun 2, 2025

[llama4] use grouped_mm in moe for sm90

3c3783b

ghstack-source-id: ffd0450 Pull Request resolved: #2755

		return permuted_indices, m_sizes, m_offsets.to(torch.int32)


		# Below is for testing only

		import triton.language as tl


		__all__ = ["generate_permute_indices"]

		From: \| rank 0 \| rank 1 \|
		To: \| E0 \| E1 \| E2 \| E3 \| E0 \| E1 \| E2 \| E3 \|

		@@ -0,0 +1,353 @@
		# Copyright (c) Meta Platforms, Inc. and affiliates.

[llama4] use grouped_mm in moe for sm90 #2755

Are you sure you want to change the base?

[llama4] use grouped_mm in moe for sm90 #2755

Uh oh!

Conversation

IvanKobzarev commented May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2755

✅ No Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented May 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ebsmothers left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ebsmothers left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

IvanKobzarev commented May 21, 2025 •

edited

Loading

pytorch-bot bot commented May 21, 2025 •

edited

Loading

codecov-commenter commented May 22, 2025 •

edited

Loading