fla error while using SFTTrainer

Checklist / 检查清单

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

Bug Description / Bug 描述

With casual_conv1d and flash_linear_attention, SFTTrainer throw cuda illegal memory access error after shard, which seems related to fla.

After flash_linear_attention uninstalled, SFTTrainer works fine though there are warning of performance degradation because of linear attention fallback to torch implementation.

Is there any version combination verified to work fine?

Error stack:
[INFO:swift] model_parameter_info: Qwen3_5MoeForConditionalGeneration: 35107.1819M Params (34639.6392M Trainable [98.6682%]), 0.0001M Buffers.
[INFO:swift] use_reentrant: True
[INFO:swift] The logging file will be saved in: /xx/logging.jsonl
[INFO:swift] Successfully registered post_encode hook: ['Qwen3_5MoeForConditionalGeneration'].
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[transformers] The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None}.
[RANK 0] Gradient accumulation steps mismatch: GradientAccumulationPlugin has 1, DeepSpeed config has 2. Using DeepSpeed's value.
Train: 0%| | 0/1082 [00:00<?, ?it/s][INFO:swift] use_logits_to_keep: False
[INFO:swift] router_aux_loss_coef: 0.0
[rank4]: Traceback (most recent call last):
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/cli/sft.py", line 20, in
[rank4]: sft_main()
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/train/sft.py", line 340, in sft_main
[rank4]: return SwiftSft(args).main()
[rank4]: ^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/base.py", line 52, in main
[rank4]: result = self.run()
[rank4]: ^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/ray_utils/base.py", line 168, in wrapper
[rank4]: return func(self, *args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/train/sft.py", line 184, in run
[rank4]: return self.train(trainer)
[rank4]: ^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/train/sft.py", line 258, in train
[rank4]: trainer.train(resume_checkpoint)
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/trainers/mixin.py", line 963, in train
[rank4]: res = super().train(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1433, in train
[rank4]: return inner_training_loop(
[rank4]: ^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1515, in _inner_training_loop
[rank4]: self._run_epoch(
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1743, in _run_epoch
[rank4]: tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/trainers/seq2seq_trainer.py", line 231, in training_step
[rank4]: return super().training_step(model, inputs, *args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1943, in training_step
[rank4]: self.accelerator.backward(loss, **kwargs)
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/accelerate/accelerator.py", line 2830, in backward
[rank4]: self.deepspeed_engine_wrapped.backward(loss, sync_gradients=self.sync_gradients, **kwargs)
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/accelerate/utils/deepspeed.py", line 270, in backward
[rank4]: self.engine.backward(loss, **kwargs)
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 33, in wrapped_fn
[rank4]: ret_val = func(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2753, in backward
[rank4]: loss.backward(**backward_kwargs)
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/_tensor.py", line 625, in backward
[rank4]: torch.autograd.backward(
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/init.py", line 354, in backward
[rank4]: _engine_run_backward(
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/graph.py", line 841, in _engine_run_backward
[rank4]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/function.py", line 315, in apply
[rank4]: return user_fn(self, *args)
[rank4]: ^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/utils/checkpoint.py", line 319, in backward
[rank4]: torch.autograd.backward(outputs_with_grad, args_with_grad)
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/init.py", line 354, in backward
[rank4]: _engine_run_backward(
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/graph.py", line 841, in _engine_run_backward
[rank4]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/function.py", line 315, in apply
[rank4]: return user_fn(self, *args)
[rank4]: ^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/utils.py", line 166, in wrapper
[rank4]: return fn(*contiguous_args, **contiguous_kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 573, in decorate_bwd
[rank4]: return bwd(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/gated_delta_rule/chunk.py", line 197, in backward
[rank4]: dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/gated_delta_rule/chunk.py", line 145, in chunk_gated_delta_rule_bwd
[rank4]: dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/utils.py", line 166, in wrapper
[rank4]: return fn(*contiguous_args, **contiguous_kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/utils/cumsum.py", line 443, in chunk_local_cumsum
[rank4]: return chunk_local_cumsum_scalar(
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/utils/cumsum.py", line 268, in chunk_local_cumsum_scalar
[rank4]: chunk_local_cumsum_scalar_kernel[grid](
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/jit.py", line 419, in
[rank4]: return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 452, in run
[rank4]: return self.fn.run(*args, **kwargs)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 236, in run
[rank4]: used_cached_result = self.check_disk_cache(key, pruned_configs, benchmark)
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 200, in check_disk_cache
[rank4]: bench_fn()
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 227, in benchmark
[rank4]: timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 162, in _bench
[rank4]: return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/testing.py", line 150, in do_bench
[rank4]: di.synchronize()
[rank4]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/cuda/init.py", line 1083, in synchronize
[rank4]: return torch._C._cuda_synchronize()
[rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank4]: torch.AcceleratorError: CUDA error: an illegal memory access was encountered
[rank4]: Search for cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information. [rank4]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. [rank4]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1 [rank4]: Compile with TORCH_USE_CUDA_DSA` to enable device-side assertions.

[rank5]: Traceback (most recent call last):
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/cli/sft.py", line 20, in
[rank5]: sft_main()
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/train/sft.py", line 340, in sft_main
[rank5]: return SwiftSft(args).main()
[rank5]: ^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/base.py", line 52, in main
[rank5]: result = self.run()
[rank5]: ^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/ray_utils/base.py", line 168, in wrapper
[rank5]: return func(self, *args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/train/sft.py", line 184, in run
[rank5]: return self.train(trainer)
[rank5]: ^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/pipelines/train/sft.py", line 258, in train
[rank5]: trainer.train(resume_checkpoint)
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/trainers/mixin.py", line 963, in train
[rank5]: res = super().train(*args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1433, in train
[rank5]: return inner_training_loop(
[rank5]: ^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1515, in _inner_training_loop
[rank5]: self._run_epoch(
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1743, in _run_epoch
[rank5]: tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/swift/trainers/seq2seq_trainer.py", line 231, in training_step
[rank5]: return super().training_step(model, inputs, *args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/transformers/trainer.py", line 1943, in training_step
[rank5]: self.accelerator.backward(loss, **kwargs)
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/accelerate/accelerator.py", line 2830, in backward
[rank5]: self.deepspeed_engine_wrapped.backward(loss, sync_gradients=self.sync_gradients, **kwargs)
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/accelerate/utils/deepspeed.py", line 270, in backward
[rank5]: self.engine.backward(loss, **kwargs)
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 33, in wrapped_fn
[rank5]: ret_val = func(*args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2753, in backward
[rank5]: loss.backward(**backward_kwargs)
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/_tensor.py", line 625, in backward
[rank5]: torch.autograd.backward(
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/init.py", line 354, in backward
[rank5]: _engine_run_backward(
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/graph.py", line 841, in _engine_run_backward
[rank5]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/function.py", line 315, in apply
[rank5]: return user_fn(self, *args)
[rank5]: ^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/utils/checkpoint.py", line 319, in backward
[rank5]: torch.autograd.backward(outputs_with_grad, args_with_grad)
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/init.py", line 354, in backward
[rank5]: _engine_run_backward(
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/graph.py", line 841, in _engine_run_backward
[rank5]: return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/autograd/function.py", line 315, in apply
[rank5]: return user_fn(self, *args)
[rank5]: ^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/utils.py", line 166, in wrapper
[rank5]: return fn(*contiguous_args, **contiguous_kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 573, in decorate_bwd
[rank5]: return bwd(*args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/gated_delta_rule/chunk.py", line 197, in backward
[rank5]: dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/gated_delta_rule/chunk.py", line 145, in chunk_gated_delta_rule_bwd
[rank5]: dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/utils.py", line 166, in wrapper
[rank5]: return fn(*contiguous_args, **contiguous_kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/utils/cumsum.py", line 443, in chunk_local_cumsum
[rank5]: return chunk_local_cumsum_scalar(
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/fla/ops/utils/cumsum.py", line 268, in chunk_local_cumsum_scalar
[rank5]: chunk_local_cumsum_scalar_kernel[grid](
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/jit.py", line 419, in
[rank5]: return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 452, in run
[rank5]: return self.fn.run(*args, **kwargs)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 236, in run
[rank5]: used_cached_result = self.check_disk_cache(key, pruned_configs, benchmark)
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 200, in check_disk_cache
[rank5]: bench_fn()
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 227, in benchmark
[rank5]: timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/runtime/autotuner.py", line 162, in _bench
[rank5]: return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/triton/testing.py", line 150, in do_bench
[rank5]: di.synchronize()
[rank5]: File "/home/xx/conda/envs/codelab/lib/python3.12/site-packages/torch/cuda/init.py", line 1083, in synchronize
[rank5]: return torch._C._cuda_synchronize()
[rank5]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank5]: torch.AcceleratorError: CUDA error: an illegal memory access was encountered
[rank5]: Search for cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information. [rank5]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect. [rank5]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1 [rank5]: Compile with TORCH_USE_CUDA_DSA` to enable device-side assertions.

[rank4]:[W629 21:46:47.312822853 ProcessGroupNCCL.cpp:1524] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
[rank5]:[W629 21:46:49.661623302 AllocatorConfig.cpp:28] Warning: PYTORCH_CUDA_ALLOC_CONF is deprecated, use PYTORCH_ALLOC_CONF instead (function operator())
[rank4]:[W629 21:46:49.661646406 AllocatorConfig.cpp:28] Warning: PYTORCH_CUDA_ALLOC_CONF is deprecated, use PYTORCH_ALLOC_CONF instead (function operator())

How to Reproduce / 如何复现

cuda: 12.8
torch: 2.9.1
python: 3.12
ms-swift: 4.4.0dev / 4.3.2
transformers: 5.12.1
triton: 3.5.1
casual_conv1d: 1.6.0 / 1.6.1 / 1.6.2.post1
flash_linear_attention: 0.4.0 / 0.4.1 (all installed from source)
fla_core: 0.4.0 / 0.4.1 (same as above)

Additional Information / 补充信息

No response

Start by reproducing SFTTrainer with casual_conv1d and flash_linear_attention, comparing it with the reported fallback after uninstalling flash_linear_attention. Trace the failure through swift/trainers/seq2seq_trainer.py into fla/ops/gated_delta_rule/chunk.py and fla/ops/utils/cumsum.py. Done means identifying a verified compatible version combination or a reproducible incompatibility.