[BUG] DeepSeek-V3 MLA rank changes trigger 16-byte data_ptr alignment failure in grouped_mm under DeepSpeed #8276

Description

DeepSeek-V3 training through DeepSpeed can fail in the Hugging Face MoE grouped_mm path with:

RuntimeError: expected data_ptr to be aligned to 16 bytes

The failure is triggered by otherwise executable MLA LoRA-rank configurations. The same frozen configurations complete under a PyTorch-native runner and a Transformers/Accelerate runner, while the DeepSpeed-wrapped runner fails consistently.

I observed this pattern for 10 configurations (5 kv_lora_rank mutations and 5 q_lora_rank mutations), and re-ran one representative case (kv_lora_rank=68) three times; it failed with the same error 3/3 times.

I am filing this in DeepSpeed because the issue is specific to the DeepSpeed execution path in our cross-runner comparison, although the immediate failing call is in Transformers/PyTorch grouped_mm, so the root cause may be an interaction between DeepSpeed and the HF MoE path.

Environment

• DeepSpeed: v0.19.1 (510ebe58e4e4148bbf49350ad18fe291a0afb381)
• Transformers: v5.9.0 (0a2757da521a7a49b8143d9e0c938f08747d682e)
• PyTorch: 2.13.0+cu126
• CUDA runtime reported by PyTorch: 12.6
• GPUs: 2x NVIDIA RTX A6000
• Python: 3.11

Representative configuration

Small DeepSeek-V3-like configuration used for the reproducer:

hidden_size=512
ffn_hidden_size=1024
num_layers=4
num_attention_heads=8
num_query_groups=2
seq_length=128
vocab_size=2048
num_experts=4
moe_router_topk=2
q_lora_rank=128
kv_lora_rank=68
qk_nope_head_dim=32
qk_rope_head_dim=32
v_head_dim=64
micro_batch_size=1
global_batch_size=4
bf16=true
fp16=false

The only mutation from the prequalified baseline in this representative case is kv_lora_rank: 64 -> 68.

Observed behavior

The failure occurs during forward in the MoE experts path:

transformers/models/deepseek_v3/modeling_deepseek_v3.py:245
  hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape)

transformers/integrations/moe.py:441
  proj_out = _grouped_linear(...)

transformers/integrations/moe.py:369
  out = _grouped_mm(input, weight.transpose(-2, -1), offs=offs)

transformers/integrations/moe.py:333
  return torch.nn.functional.grouped_mm(input.to(weight.dtype), weight, offs=offs)

torch/nn/functional.py:7139
  return torch._grouped_mm(...)

RuntimeError: expected data_ptr to be aligned to 16 bytes

Both distributed ranks fail with the same message.

Cross-runner comparison

For the same frozen mutation intents:

• PyTorch-native: completes
• Transformers/Accelerate: completes
• DeepSpeed: fails with the 16-byte data_ptr alignment error

This was observed for:

• kv_lora_rank: 60, 65, 66, 67, 68
• q_lora_rank: 124, 129, 130, 131, 132

Expected behavior

Wrapping the model with DeepSpeed should not make these otherwise executable shapes produce a misaligned tensor passed into grouped_mm, or the integration should normalize/copy the relevant tensor before the kernel call if a stronger alignment requirement is introduced.