Megatron Train Qwen3.5-VL-35B-A3B Error

Checklist / 检查清单

• I have searched existing issues, and this is a new bug report. / 我已经搜索过现有的 issues，确认这是一个新的 bug report。

Bug Description / Bug 描述

transformer_engine==2.12.0
transformer_engine_cu12==2.12.0
transformer_engine_torch==2.12.0
transformers==5.2.0
transformers-stream-generator==0.0.5
megatron-core==0.15.3
flash_attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp312-cp312-linux_x86_64.whl#sha256=f25da18657a87fc83dc1bfb8b7751b82246e9db355510226b674fd437c34b5fb

Editable Git install with no remote (ms_swift==4.1.0.dev0)

-e /opt/tiger/ms-swift
torchaudio==2.8.0+cu129
torchvision==0.23.0+cu129

基于上述 Python 环境训练 By Megatron 报错信息:
[rank2]: Traceback (most recent call last):
[rank2]: File "/opt/tiger/ms-swift/swift/cli/_megatron/sft.py", line 7, in
[rank2]: megatron_sft_main()
[rank2]: File "/opt/tiger/ms-swift/swift/megatron/pipelines/train/sft.py", line 88, in megatron_sft_main
[rank2]: return MegatronSft(args).main()
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/opt/tiger/ms-swift/swift/pipelines/base.py", line 52, in main
[rank2]: result = self.run()
[rank2]: ^^^^^^^^^^
[rank2]: File "/opt/tiger/ms-swift/swift/megatron/pipelines/train/sft.py", line 65, in run
[rank2]: trainer = self.prepare_trainer()
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/opt/tiger/ms-swift/swift/megatron/pipelines/train/sft.py", line 33, in prepare_trainer
[rank2]: return MegatronTrainer(self.args, self.template)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/opt/tiger/ms-swift/swift/megatron/trainers/base.py", line 95, in init
[rank2]: warmup_jit_function(self.config, args)
[rank2]: File "/opt/tiger/ms-swift/swift/megatron/utils/megatron_lm_utils.py", line 681, in warmup_jit_function
[rank2]: output = bias_swiglu(input_tensor, bias)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 749, in compile_wrapper
[rank2]: raise e.remove_dynamo_frames() from None # see TORCHDYNAMO_VERBOSE=1
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 1871, in _call_user_compiler
[rank2]: raise BackendCompilerFailed(
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 1846, in call_user_compiler
[rank2]: compiled_fn = compiler_fn(gm, example_inputs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/dynamo/repro/after_dynamo.py", line 150, in call
[rank2]: compiled_gm = compiler_fn(gm, example_inputs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/init.py", line 2380, in call
[rank2]: return compile_fx(model, inputs, config_patches=self.config)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2418, in compile_fx
[rank2]: return aot_autograd(
[rank2]: ^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/backends/common.py", line 109, in call
[rank2]: cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1199, in aot_module_simplified
[rank2]: compiled_fn = AOTAutogradCache.load(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/autograd_cache.py", line 1140, in load
[rank2]: compiled_fn = dispatch_and_compile()
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1184, in dispatch_and_compile
[rank2]: compiled_fn, _ = create_aot_dispatcher_function(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 576, in create_aot_dispatcher_function
[rank2]: return _create_aot_dispatcher_function(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 836, in _create_aot_dispatcher_function
[rank2]: compiled_fn, fw_metadata = compiler_fn(
[rank2]: ^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 1605, in aot_dispatch_autograd
[rank2]: compiled_fw_func = aot_config.fw_compiler(fw_module, adjusted_flat_args)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 483, in call
[rank2]: return self.compiler_fn(gm, example_inputs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 2250, in fw_compiler_base
[rank2]: return inner_compile(
[rank2]: ^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 745, in compile_fx_inner
[rank2]: return wrap_compiler_debug(_compile_fx_inner, compiler_name="inductor")(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_dynamo/repro/after_aot.py", line 124, in debug_wrapper
[rank2]: inner_compiled_fn = compiler_fn(gm, example_inputs)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 860, in _compile_fx_inner
[rank2]: (key_info, cache_info) = FxGraphCache.prepare_key(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/codecache.py", line 1474, in prepare_key
[rank2]: key, debug_lines = compiled_fx_graph_hash(
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/codecache.py", line 960, in compiled_fx_graph_hash
[rank2]: details = FxGraphHashDetails(gm, example_inputs, fx_kwargs, inputs_to_check)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/codecache.py", line 896, in init
[rank2]: self.system_info = CacheBase.get_system()
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/usr/local/lib/python3.12/site-packages/torch/_inductor/codecache.py", line 205, in get_system
[rank2]: from triton.compiler.compiler import triton_key
[rank2]: torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
[rank2]: ImportError: cannot import name 'triton_key' from 'triton.compiler.compiler' (/usr/local/lib/python3.12/site-packages/triton/compiler/compiler.py)

[rank2]: Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

How to Reproduce / 如何复现

训练脚本如下: PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True'
NPROC_PER_NODE=4
CUDA_VISIBLE_DEVICES=0,1,2,3
megatron sft
--model $model_path
--save_safetensors true
--merge_lora true
--dataset 'AI-ModelScope/LongAlpaca-12k'
--load_from_cache_file true
--add_non_thinking_prefix true
--split_dataset_ratio 0.01
--tuner_type lora
--lora_rank 8
--lora_alpha 32
--target_modules all-linear
--expert_model_parallel_size 4
--moe_permute_fusion true
--moe_grouped_gemm true
--moe_shared_expert_overlap true
--moe_aux_loss_coeff 1e-6
--micro_batch_size 1
--global_batch_size 16
--recompute_granularity full
--recompute_method uniform
--recompute_num_layers 1
--num_train_epochs 1
--group_by_length true
--finetune true
--freeze_llm false
--freeze_vit true
--freeze_aligner true
--cross_entropy_loss_fusion true
--lr 1e-4
--lr_warmup_fraction 0.05
--min_lr 1e-5
--output_dir $output_path
--eval_steps 200
--save_steps 200
--max_length 8192
--dataloader_num_workers 8
--dataset_num_proc 8
--no_save_optim true
--no_save_rng true
--sequence_parallel true
--attention_backend flash
--padding_free false

Additional Information / 补充信息

辛苦帮忙看下

Start at swift/megatron/utils/megatron_lm_utils.py in warmup_jit_function, where the traceback reaches bias_swiglu, and review the listed PyTorch, Triton, and Megatron-related versions. Reproduce with the provided four-GPU megatron sft command and verify that the warmup completes without the triton_key import error and training proceeds.