Deepseekv3 CUDA illegal memory access on B200 from increasing seq_len to 8192

[danielvegamyhre]

Bug description

Using B200 devgpu

Repro

• fsdp=4, ep=4, AC=none, compile=False

• Using debug model with DSV3 671b config in single node testing via setting n_layers=2 and n_dense_layers=1 to avoid OOM.

  • Is this not a valid config? Could this explain the error?

    "debugmodel": DeepSeekV3ModelArgs(
        vocab_size=129280,
        dim=7168,
        inter_dim=18432,
        moe_inter_dim=2048,
        n_layers=2,
        n_dense_layers=1,
        n_heads=128,
        moe_args=MoEArgs(
            num_experts=256,
            num_shared_experts=1,
            top_k=8,
            score_func="sigmoid",
            route_norm=True,
            route_scale=2.5,
            score_before_experts=False,
        ),
        n_expert_groups=8,
        n_limited_groups=4,
        q_lora_rank=1536,
        kv_lora_rank=512,
        qk_nope_head_dim=128,
        qk_rope_head_dim=64,
        v_head_dim=128,
        use_flex_attn=True,
        attn_mask_type="block_causal",
    ),

Default seq len (2048) works:

CUDA_VISIBLE_DEVICES="4,5,6,7" NGPU=4 CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ./run_train.sh --training.steps=200 --parallelism.data_parallel_shard_degree=4 --parallelism.expert_parallel_degree=4 --parallelism.tensor_parallel_degree=1 --activation_checkpoint.mode=none 

Increasing seq_len to 8192 -> CUDA illegal memory access:

CUDA_VISIBLE_DEVICES="4,5,6,7" NGPU=4 CONFIG_FILE="./torchtitan/models/deepseek_v3/train_configs/debug_model.toml" ./run_train.sh --training.steps=200 --parallelism.data_parallel_shard_degree=4 --parallelism.expert_parallel_degree=4 --parallelism.tensor_parallel_degree=1 --activation_checkpoint.mode=none --training.seq_len=8192

Versions

• torchtitan latest main branch
• torch 10/2 nightly for cuda 12.8: torch 2.10.0.dev20251002+cu128

[danielvegamyhre]

This repros with deepseek_v3_16b.toml with default seq len:

CONFIG_FILE=/home/danvm/torchtitan/torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml NGPU=8 ./run_train.sh --metrics.log_freq=10 --training.steps=200  --parallelism.data_parallel_shard_degree=8 --parallelism.expert_parallel_degree=8 --parallelism.tensor_parallel_degree=1 --compile.enable --training.seq_len=2048 --activation_checkpoint.mode=none --model.print_after_conversion --profiling.enable_profiling --profiling.profile_freq=50

[danielvegamyhre]

Running with removing --log-filter-ranks 0 and setting CUDA_LAUNCH_BLOCKING=1 it looks like for triton_flex_attention there are both:

• out of memory
• IMAs

[rank7]:[rank7]:E1003 09:48:28.513000 2742466 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] No valid triton configs. OutOfMemoryError: out of resource: triton_flex_attention_backward Required: 298816 Hardware limit:232448 Reducing block sizes or `num_stages` may help.. 
[rank7]:[rank7]:E1003 09:48:28.513000 2742466 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] Ignoring this choice.
[rank7]:[rank7]:E1003 09:48:28.529000 2742466 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] Runtime error during autotuning: 
[rank7]:[rank7]:E1003 09:48:28.529000 2742466 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] No valid triton configs. OutOfMemoryError: out of resource: triton_flex_attention_backward Required: 348224 Hardware limit:232448 Reducing block sizes or `num_stages` may help.. 
[rank7]:[rank7]:E1003 09:48:28.529000 2742466 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] Ignoring this choice.
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] Runtime error during autotuning: 
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] CUDA error: unspecified launch failure
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] Search for `cudaErrorLaunchFailure' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] For debugging consider passing CUDA_LAUNCH_BLOCKING=1
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
[rank4]:[rank4]:E1003 09:48:28.951000 2742461 site-packages/torch/_inductor/select_algorithm.py:3213] [1/0] . 

Full logs: https://www.internalfb.com/intern/everpaste/?handle=GMsaEiECcYyt98oCAC0_XPUsUURPbsIXAAAz&phabricator_paste_number=1975644367

@drisspg is flex attention supposed to support b200/sm100 yet?

[drisspg]

Yeah, can you compile flex-attention = max-autotune-no-cudagraphs

[fegin]

@drisspg TorchTitan uses max-autotune-no-cudagraphs.

[danielvegamyhre]

I set use_flex_attn=False as a workaround for now, which has worked.

[drisspg]

If you could run computer sanitizer that would be pretty helpful

[syed-ahmed]

@danielvegamyhre Were you able triage this further or should we take a look at this?

[danielvegamyhre]

@danielvegamyhre Were you able triage this further or should we take a look at this?

No I just stuck with my workaround of not using flex attention so far, if you could help take a look that would be great!

[syed-ahmed]

Just a quick update:

• I can reproduce this standalone using attention-gym with sequence length 4096 and head dimension 192, on b200. Config file changed to:

diff --git a/benchmarks/configs/config_basic.yaml b/benchmarks/configs/config_basic.yaml
index 53aae05..1ad1b1c 100644
--- a/benchmarks/configs/config_basic.yaml
+++ b/benchmarks/configs/config_basic.yaml
@@ -7,10 +7,10 @@ calculate_bwd: false
 dtype: "bfloat16"
 
 # Shape parameters
-b: [2, 4]  # batch sizes
+b: [4]  # batch sizes
 nh: ["16, 16", "16, 4"]  # [query_heads,key_value_heads]
-s: [1024, 2048, 4096, 8192, 16384, 32768]  # sequence lengths
-d: [64, 128]  # head dimensions
+s: [4096]  # sequence lengths
+d: [192]  # head dimensions

• The torch.AcceleratorError: CUDA error: unspecified launch failure is the main error that doesn't let the workload continue. It only happens for the backward compilation. Command used: python benchmarks/flex_perf.py --config benchmarks/configs/config_basic.yaml --dtype float16 --max_autotune true --calculate_bwd true
• Core dump (using CUDA_ENABLE_COREDUMP_ON_EXCEPTION=1 CUDA_ENABLE_LIGHTWEIGHT_COREDUMP=1 CUDA_COREDUMP_SHOW_PROGRESS=1 , and then (cuda-gdb) target cudacore <coredump_name>.nvcudmp) shows the following trace, confirming the error coming from triton_flex_attention_backward:

  #0  0x0000fff576bf7e20 in triton_flex_attention_backward<<<(64,4,16),(128,1,1)>>> () at /tmp/torchinductor_root/ec/cecqmpj5xz5y6e6bapissrszpr27hw4f53ky6mismyp4rdsld5et.py:18
  18      def triton_flex_attention_backward(arg_Q, arg_K, arg_V, arg_LSE, arg_DELTA, arg_DO, arg_DQ, arg_DV, arg_KV_NUM_BLKS, arg_KV_IDX, arg_Q_NUM_BLKS, arg_Q_IDX, arg_FULL_KV_NUM_BLKS, arg_FULL_KV_IDX, arg_FULL_Q_NUM_BLKS, arg_FULL_Q_IDX, in_ptr16, out_ptr0):
  (
  

• In addition, export CUDA_LOG_FILE=cuda.err.log shows the following:

  [18:34:02.306][281435463938432][CUDA][E] Sticky error detected
  [18:34:02.306][281435463938432][CUDA][E] Returning 719 (CUDA_ERROR_LAUNCH_FAILED) from cuLaunchKernel
  [18:34:02.307][281435463938432][CUDA][E] Sticky error detected
  [18:34:02.307][281435463938432][CUDA][E] Returning 719 (CUDA_ERROR_LAUNCH_FAILED) from cuLaunchKernel
  [18:34:02.307][281435463938432][CUDA][E] Sticky error detected
  

I'll triage this further with our team. My hunch is that either the consecutive OOMs during autotuning is corrupting something, or we actually need to update the heuristics, similar to this PR: pytorch/pytorch#160192.

[drisspg]

Im glad you were able to debug some more @syed-ahmed I was also able to repro this standalone, seems pretty isolated tot he head_dimqk =192 aand head_dimv = 128 case

[Aidyn-A]

select_algorithm.py:3213 [1/0] No valid triton configs. OutOfMemoryError: out of resource: triton_flex_attention_backward >Required: 298816 Hardware limit:232448 Reducing block sizes or `num_stages` may help.. 
select_algorithm.py:3213 [1/0] Ignoring this choice.
select_algorithm.py:3213 [1/0] Runtime error during autotuning: 
select_algorithm.py:3213 [1/0] No valid triton configs. OutOfMemoryError: out of resource: triton_flex_attention_backward >Required: 348224 Hardware limit:232448 Reducing block sizes or `num_stages` may help.. 

I have encountered these with some extreme/unsuitable for Blackwell configs.