issues on llama3 compile + (async) TP + AC

[tianyu-l]

Bug description

on 8 GPUs, DP2 TP4

1. compile + selective op AC + TP:
   got failure

File "/data/users/lty/pytorch/torch/_ops.py", line 1317, in __getattr__
      raise AttributeError(
  AttributeError: '_OpNamespace' 'symm_mem' object has no attribute 'fused_all_gather_matmul'

Note: DP4 TP2 works.

2. compile + selective op AC + async TP:
   got very low throughput, compared with "full AC + TP"

3. compile + full AC + asyncTP:
   got the following warning and very low throughput (compared with "full AC + TP")

torch/_inductor/fx_passes/micro_pipeline_tp.py:894] [0/1] no producer matmul found for reduce scatter, skipping fuse_matmul_reduce_scatter fusion

4. compile + async TP (+ selective 2 ac) is still failing on CI machines https://github.com/pytorch/torchtitan/actions/runs/14992456541/job/42118899398?pr=1186

Versions

latest pytorch built from source
on A100 GPUs
debug model

[tianyu-l]

cc @kwen2501 @danielvegamyhre @bdhirsh

[danielvegamyhre]

AttributeError: '_OpNamespace' 'symm_mem' object has no attribute 'fused_all_gather_matmul'

Hmm, I've never seen this before, seems like the custom op isn't registered somehow?

torch/_inductor/fx_passes/micro_pipeline_tp.py:894] [0/1] no producer matmul found for reduce scatter, skipping fuse_matmul_reduce_scatter fusion

This warning is expected, many valid graphs will not have candidates for matmul reduce scatter fusion. If I remember correctly the backward pass is often like this because we are adding gradients then doing reduce-scatter, which doesn't match the matmul reduce scatter pattern.

compile + async TP (+ selective 2 ac) is still failing on CI machines https://github.com/pytorch/torchtitan/actions/runs/14992456541/job/42118899398?pr=1186

This looks like a CUDA driver error on the host machine.

    File "/opt/conda/envs/py_3.11/lib/python3.11/site-packages/torch/distributed/_symmetric_memory/__init__.py", line 120, in get_symm_mem_workspace
      return _SymmetricMemory.rendezvous(tensor)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  RuntimeError: CUDA driver error: invalid device ordinal

[fegin]

on 8 GPUs, DP2 TP4
compile + selective op AC + TP:
got failure

@tianyu-l What's your command? I could not reproduce the same issue on A100 nor on H100. This is my command:

CONFIG_FILE="torchtitan/models/llama3/train_configs/llama3_8b.toml " ./run_train.sh  --training.compile --parallelism.tensor_parallel_degree 4 --parallelism.enable_async_tensor_parallel

As for the performance issue, like @danielvegamyhre mentioned this is expected due to the "noise" nodes in the graph. Luca has mentioned some ideas to improve async TP, which we will explore next.

The CI error seems to be a transient infra/machine error as it is green and the original error seems to be related to the infra/machine setup error.

[tianyu-l]

@fegin

This is my command:

You also need selective_ac_option = 'op' instead of the default 2

The CI error seems to be a transient infra/machine error as it is green and the original error seems to be related to the infra/machine setup error.

Async TP CI has been commented out for months. Even if it's related to CI machine, we should still fix.

[fegin]

@tianyu-l

You also need selective_ac_option = 'op' instead of the default 2

That's the default value. 2 is not default. And I verified that even if I specify --activation_checkpoint.selective_ac_option="op", it still works.

[tianyu-l]

@fegin
After I switch to another H100 machine, it works for me.

That's the default value. 2 is not default.

oh I ran on the debugmodel, where default is 2.
https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/llama3/train_configs/debug_model.toml#L68

[danielvegamyhre]

@tianyu-l I couldn't repro this error, using either per op SAC or per layer SAC. I suspect it is a pytorch build issue, since that custom op definitely exists. I would try a fresh re-install w/ 2.6.0+ (not sure oldest torch version that would include these ops?) or fresh build from source.

File "/data/users/lty/pytorch/torch/_ops.py", line 1317, in __getattr__
      raise AttributeError(
  AttributeError: '_OpNamespace' 'symm_mem' object has no attribute 'fused_all_gather_matmul'

[danielvegamyhre]

Update: @tianyu-l I actually accidentally reproed the no attribute 'fused_all_gather_matmul' error, strangely not with async TP but with vanilla TP. To me this indicates it is an inductor caching issue, since the original dynamo traced graph (before micro_pipeline_tp post grad pass) doesn't change between async TP and vanilla TP, so it sees a cached inductor graph/codegen which matches the trace and tries to run it. However, the cached inductor codegen is the product of a post-grad graph manipulation where symmetric memory was used, and now with vanilla tp the symmetric memory group is not initialized so the custom op is not registered.

Clearing the inductor cache and rerunning resolved the issue.

rm -rf /tmp/torchinductor_${USER}; CONFIG_FILE="torchtitan/models/llama3/train_
configs/llama3_8b.toml " ./run_train.sh  --training.compile --parallelism.tensor_parallel_degree 4

I can see this potentially confusing users though, maybe we should think about how to ensure this inductor cache from async TP isn't used for TP? Maybe @bdhirsh has thoughts on how to do this?

[fegin]

@danielvegamyhre Very interesting finding. I thought inductor would prefix the timestamp to the cache folder. But it looks like I'm wrong.