Evo2 FP8 TransformerEngine Linear hits CUDA illegal memory access under 4-GPU DeepSpeed on H100

[ShiLong-CN]

Summary

We are seeing a reproducible CUDA illegal memory access when running Evo2 with FP8 enabled under 4-GPU DeepSpeed training on H100.

The same setup works in two control cases:

• Single-GPU Evo2 FP8: works
• 4-GPU DeepSpeed with EVO2_DISABLE_FP8=1: works
• 4-GPU DeepSpeed with Evo2 FP8 enabled: fails during the first training step

This suggests the issue is specific to Evo2 FP8 / TransformerEngine execution in distributed multi-GPU training.

Environment

• GPU: NVIDIA H100 80GB HBM3
• Python: 3.11.15
• PyTorch: 2.7.1+cu128
• CUDA: 12.8
• evo2: 0.5.5
• transformer_engine: 2.3.0
• transformer_engine_torch: 2.3.0
• flash-attn: 2.8.0.post2
• deepspeed: 0.14.4
• lightning: 2.3.3

Note: Evo2 prints the following warning at import time:

Supported flash-attn versions are >= 2.1.1, <= 2.7.4.post1. Found flash-attn 2.8.0.post2.

Reproduction

The failing run uses 4 H100 GPUs with DeepSpeed stage 2:

CUDA_VISIBLE_DEVICES=4,5,6,7 \
python -m torch.distributed.run \
  --standalone \
  --nnodes=1 \
  --nproc_per_node=4 \
  train.py \
  ... \
  --devices 4 \
  --strategy deepspeed_stage_2 \
  --max_steps 1

EVO2_DISABLE_FP8 is unset.

The run successfully loads the models, initializes DeepSpeed/NCCL, and enters training, then fails on the first step.

Observed Error

The root error happens inside TransformerEngine GEMM during Evo2 forward:

Training:   0%|          | 0/8804

File ".../train.py", line 532, in training_step
File ".../models/dna_llm.py", line 789, in forward
File ".../evo2/models.py", line 103, in forward
File ".../vortex/model/model.py", line 490, in proj_norm
File ".../vortex/model/layers.py", line 82, in forward
File ".../transformer_engine/pytorch/module/linear.py", line 1289, in forward
File ".../transformer_engine/pytorch/module/linear.py", line 276, in forward
File ".../transformer_engine/pytorch/cpp_extensions/gemm.py", line 111, in general_gemm

RuntimeError: CUDA error: an illegal memory access was encountered

The same error appears across multiple ranks; NCCL cleanup errors appear afterward and seem secondary.

Control Runs

Single-GPU FP8 works

CUDA_VISIBLE_DEVICES=4
# EVO2_DISABLE_FP8 unset
# max_steps=1

Result:

train/loss_step=2.840
Trainer.fit stopped: max_steps=1 reached

4-GPU with FP8 disabled works

CUDA_VISIBLE_DEVICES=0,1,2,3
EVO2_DISABLE_FP8=1
# max_steps=1

Result:

train/loss_step=3.250
Trainer.fit stopped: max_steps=1 reached

Expected Behavior

Evo2 FP8 forward should either work under multi-GPU DeepSpeed training, or the documentation should specify the supported distributed-training configuration and recommended package versions.

Workaround

Setting the following environment variable avoids the crash:

EVO2_DISABLE_FP8=1

This allows 4-GPU DeepSpeed training to proceed, but disables Evo2 FP8.

Questions

1. Is Evo2 FP8 currently expected to support multi-GPU distributed training with DeepSpeed?
2. Is there a recommended version combination for torch, transformer_engine, flash-attn, and CUDA on H100?
3. Could the unsupported flash-attn==2.8.0.post2 version plausibly trigger this TransformerEngine FP8 GEMM failure, or is this likely independent?

[garykbrixi]

Hi @ShiLong-CN, this evo 2 repo is for inference not for training. To train Evo 2, please use Savanna or BioNemo which provides a Evo 2 finetuning tutorial here.

[ShiLong-CN]

Hi @garykbrixi, thank you for the clarification.

Just to clarify our setup: we are not trying to fine-tune Evo2 itself. Evo2 is kept frozen and is only used as an inference-only DNA encoder inside our larger Lightning/DeepSpeed model. The failing call is Evo2(..., return_embeddings=True, layer_names=...), used to extract hidden states for downstream cross-attention.

The crash happens before any Evo2 backward/update. In our 4-GPU run, each distributed rank launches its own process and loads its own Evo2 instance. Each rank then calls frozen Evo2 forward on its local batch. FP8 input projections are enabled on H100.

Our current hypothesis is that this may not be an "Evo2 training" issue, but a multi-process FP8 inference issue: four ranks each initialize their own Evo2 / TransformerEngine FP8 state, while our wrapper also moves the Evo2 model to the rank-local device and updates block_idx_to_device. The failure occurs inside the frozen Evo2 forward path, around:

Evo2.forward -> vortex/model/model.py::proj_norm -> TransformerEngine Linear -> fp8_gemm / cuBLASLt

We have observed CUDA illegal memory access and cuBLASLt initialization errors there.

Control runs:

• single GPU + Evo2 FP8 works
• 4 GPU + EVO2_DISABLE_FP8=1 works
• 4 GPU + Evo2 FP8 enabled fails

Do you expect Evo2 FP8 inference to support this kind of multi-process data-parallel usage on one H100 node, where each rank owns a separate Evo2 instance? Or is the supported path only single-process inference / Vortex-managed model parallelism?

Also, should user code avoid manually calling .to(local_rank_device) or modifying block_idx_to_device when using Evo2 FP8?

[garykbrixi]

Thanks @ShiLong-CN for more information. The one workaround for your setup may be to use Evo 2 on BioNemo. Multi-process data-parallel FP8 inference is not a tested configuration here - the FP8 path was designed for single-process inference with vortex-managed device placement due to the sensitivity of TransformerEngine's linears.

We noticed during development that manually calling .to(local_rank_device) can cause issues for the fp8 linear layers from TransformerEngine, and we apply a fixup inside vortex to deal with this. I suspect the device placement and movement you are describing could be causing the TransformerEngine linear layer extra state to be on the wrong device, as we were seeing with @antonvnv -- see PR here for more context