autobucketing does not support deepseekv3

[BoyuanFeng]

Repro:

NGPU=4 CONFIG_FILE=/data/users/boyuan/torchtitan/torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml ./run_train.sh --model.name compiler_toolkit.deepseek_v3 --parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=2 --parallelism.expert_parallel_degree=2 --activation_checkpoint.mode none --job.custom_config_module=torchtitan.experiments.compiler_toolkit.job_config --compile.passes autobucketing_reordering

Error:

    File "/data/users/boyuan/torchtitan/torchtitan/experiments/compiler_toolkit/passes.py", line 28, in autobucketing_reordering_pass
      schedule_overlap_bucketing(gm, collective_bucketing=True, collective_estimator = "benchmark",)
    File "/data/users/boyuan/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 1032, in schedule_overlap_bucketing
      return OverlapScheduler(
             ^^^^^^^^^^^^^^^^^
    File "/data/users/boyuan/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 281, in __init__
      self._identify_collectives()
    File "/data/users/boyuan/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 323, in _identify_collectives
      coll_time_ms = estimate_collective_time(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/data/users/boyuan/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 65, in estimate_collective_time
      return torch._inductor.comm_analysis.estimate_nccl_collective_runtime_from_fx_node(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/data/users/boyuan/pytorch/torch/_inductor/comm_analysis.py", line 397, in estimate_nccl_collective_runtime_from_fx_node
      coll = get_collective_type_from_kernel_name(fx_node.target.name())
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/data/users/boyuan/pytorch/torch/_inductor/comm_analysis.py", line 59, in get_collective_type_from_kernel_name
      raise ValueError(f"Unsupported collective kernel: {kernel_name}")
  ValueError: Unsupported collective kernel: _c10d_functional::all_to_all_single

I tried analytical mode for autobucketing, which errors with all_to_all not supported.

I also tried benchmark mode. It shows the same error all_to_all not supported. I checked a bit. No matter values of collective_estimator, OverlapScheduler calls _identify_collectives() anyway, which finally uses analytic model.

cc @eellison @ruisizhang123 @tianyu-l

[ruisizhang123]

fixed in: pytorch/pytorch#167797

[yiming0416]

I am getting another error when running dsv3 with autobucketing pass after including the fix. Does it ring a bell to you? @ruisizhang123

Repro:

NGPU=4 CONFIG_FILE=./torchtitan/models/deepseek_v3/train_configs/debug_model.toml ./run_train.sh --model.name simple_fsdp.deepseek_v3 --parallelism.data_parallel_shard_degree=2 --parallelism.tensor_parallel_degree=2 --parallelism.expert_parallel_degree=2 --activation_checkpoint.mode none --job.custom_config_module=torchtitan.experiments.simple_fsdp.job_config --compile.enable --compile.graph_passes auto_bucketing --compile.backend aot_eager

The error

    File "/data/users/yimingzhou/torchtitan/torchtitan/experiments/simple_fsdp/backend.py", line 55, in aot_eager_autobucketing_reordering_pass
      schedule_overlap_bucketing(gm)
    File "/data/users/yimingzhou/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 1025, in schedule_overlap_bucketing
      return OverlapScheduler(
             ^^^^^^^^^^^^^^^^^
    File "/data/users/yimingzhou/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 272, in __init__
      self._identify_collectives()
    File "/data/users/yimingzhou/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 314, in _identify_collectives
      coll_time_ms = estimate_collective_time(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/data/users/yimingzhou/pytorch/torch/_inductor/fx_passes/overlap_scheduling.py", line 66, in estimate_collective_time
      return torch._inductor.comm_analysis.estimate_nccl_collective_runtime_from_fx_node(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/data/users/yimingzhou/pytorch/torch/_inductor/comm_analysis.py", line 372, in estimate_nccl_collective_runtime_from_fx_node
      tensor_storage_size_bytes = estimate_fx_collective_size(fx_node)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/data/users/yimingzhou/pytorch/torch/_inductor/comm_analysis.py", line 347, in estimate_fx_collective_size
      numel = get_size_numel(t.size())
                             ^^^^^^
  torch._dynamo.exc.BackendCompilerFailed: backend='simple_fsdp_custom_pass' raised:
  AttributeError: 'SymInt' object has no attribute 'size'

[ruisizhang123]

ehhhh seems like dynamic shape thing, will submit a fix for it.