Skip to content

Preserve FP32 phase construction for trained RoPE#1

Merged
Bias92 merged 1 commit into
Bias92:add-exaone4_5-inference-v2from
tohtana:tohtana/pr8121-fp32-trained-rope
Jul 15, 2026
Merged

Preserve FP32 phase construction for trained RoPE#1
Bias92 merged 1 commit into
Bias92:add-exaone4_5-inference-v2from
tohtana:tohtana/pr8121-fp32-trained-rope

Conversation

@tohtana

@tohtana tohtana commented Jul 14, 2026

Copy link
Copy Markdown

Summary

  • Extend the existing trained-RoPE CUDA path so FP16/BF16 Q/K and KV cache inputs consume FP32 inverse frequencies and construct phases in FP32 before rotated outputs are cast back to the activation dtype.
  • Keep the EXAONE 4.5 inverse-frequency buffer in FP32 without changing its 4096-token scope or hybrid-attention dispatch.
  • Add FP16 and BF16 actual-kernel coverage at position 4095.

Testing

  • python -m pytest -q -s tests/unit/inference/v2/model_implementations/test_exaone4_5.py (6 passed)
  • python -m pytest -q -s -m inference_v2_ops tests/unit/inference/v2/kernels/ragged_ops/test_blocked_rotary_emb.py (58 passed)
  • python -m pytest -q -s -m inference_v2_ops tests/unit/inference/v2/modules/test_blocked_attn.py (23 passed)
  • Actual trained-RoPE CUDA parity against the Transformers float32-phase reference at position 4095 for FP16 and BF16
  • Real LGAI-EXAONE/EXAONE-4.5-33B prefill/decode smoke with the 4096-token cap and hybrid-attention dispatch checks

@tohtana
Preserve FP32 trained RoPE phase construction
73c325e 
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
@tohtana tohtana mentioned this pull request Jul 14, 2026
Merged
@Bias92
Bias92 merged commit 73c325e into Bias92:add-exaone4_5-inference-v2 Jul 15, 2026
banxingmjj pushed a commit to openanolis/DeepSpeed that referenced this pull request Jul 15, 2026
@Bias92 @tohtana
Add EXAONE 4.5 model support for Inference V2 (deepspeedai#8121)
d36079f 
# Add EXAONE 4.5 model support for Inference V2

## Summary
Add support for LG AI Research's EXAONE 4.5 language model in DeepSpeed
Inference V2.

The EXAONE 4.5 text decoder shares EXAONE 4.0's post-norm + QK-Norm
parameter layout (`text_config.architectures ==
["Exaone4ForCausalLM"]`), so the transformer container and most of the
inference model are reused. EXAONE 4.5 is also a hybrid-attention model:
`sliding_attention` layers apply llama3-scaled RoPE, and
`full_attention` layers use no positional embedding (global NoPE). This
implementation follows that reference behavior.

## Changes
- New model implementation
`deepspeed/inference/v2/model_implementations/exaone4_5/`:
- `container.py`: non-transformer container for the multimodal
checkpoint layout. The LM weights are nested under
`model.language_model.` and `lm_head` stays top-level.
- `model.py`: hybrid attention on top of the reused EXAONE 4.0 forward.
- `sliding_attention` layers use trained-frequency RoPE with
llama3-scaled inverse frequencies. A unit test checks the values against
transformers' `ROPE_INIT_FUNCTIONS["llama3"]`.
- `full_attention` layers dispatch to a separate NoPE attention module.
- Sequence length is capped at `sliding_window` (4096). The dense
blocked attention kernel has no local mask, and at or below the window
size dense causal attention is equivalent to sliding attention. The cap
is enforced in the scheduler path (`get_kv_requirements`) and the direct
path (`maybe_allocate_kv`).
- `activation_dtype` handles the transformers v5 config (`dtype` vs
`torch_dtype`, str vs `torch.dtype`).
- `policy.py`: extracts `text_config`, reuses
`Exaone4TransformerContainer` under the `model.language_model.layers`
prefix, and leaves the vision tower (`model.visual.`) and the MTP head
(`mtp.`) unmapped.
- `exaone4/model.py`: extracted a behavior-neutral `_forward_attention`
seam for the per-layer dispatch, and fixed a latent aliasing crash in
`_apply_qk_norm`. For a single-row slice, `contiguous()` returns an
alias rather than a copy, so the write-back overlapped its own source
and every single-token decode step raised a RuntimeError. The q/k slices
are now cloned. This also affected EXAONE 4.0, which had not been
exercised on the decode path.
- `checkpoint/huggingface_engine.py`: derive `max_seq_length` from the
nested `text_config.max_position_embeddings` when the top-level
(multimodal) config doesn't have it.
- `kernels/ragged_ops/linear_blocked_kv_rotary`: the trained-RoPE path
now keeps the inverse frequencies and phase construction in FP32 instead
of coupling them to the Q/K activation dtype (repair by @tohtana, merged
from [Bias92#1](Bias92#1)).
Q/K/V and the KV cache stay FP16/BF16; the rotated outputs are cast back
to the activation dtype. A kernel test at position 4095 covers both
dtypes.
- Registered `exaone4_5` in `engine_factory.py` and
`model_implementations/__init__.py`.
- Unit tests for the llama3 frequency computation (against
transformers), the per-layer attention dispatch, and the sequence-length
cap.

## Why the mapping is exact
Checked against the real `LGAI-EXAONE/EXAONE-4.5-33B` checkpoint index
(1064 tensors). All 64 decoder layers expose exactly the 11 parameters
the 4.0 container maps, and every non-language-model tensor falls under
`model.visual.` (342) or `mtp.` (15), both declared unmapped.

## Testing
Validated on an A100 80GB (SXM4) at `73c325e`, CUDA 12.4, torch
2.4.1+cu124 / transformers 5.13.1, with the real
`LGAI-EXAONE/EXAONE-4.5-33B` checkpoint (bf16, shard SHA-256 verified):

- Unit tests for the diff-impacted scope, run on CI-matched pytest
8.3.5: rotary kernels 58/58, blocked attention 23/23, EXAONE 4.5 model
6/6.
- FP32-phase kernel test at position 4095 passes for both FP16 and BF16.
- Load: both shards load and every container initializes (mapping and
the `model.visual.`/`mtp.` skip both work). `max_sequence_length`
resolves to the enforced 4096 cap and the runtime rope buffer is
`torch.float32`.
- All pre-commit hooks (including clang-format on the CUDA files) pass
on the changed files.
- Greedy generation through `engine.put` (single-sequence decode, which
exercises the trained-rotary kernel, the NoPE module and the per-layer
dispatch):

```
'The capital of France is' -> ' Paris.  \nThe capital of Germany is Berlin.  \nThe capital of'
'대한민국의 수도는'          -> ' 어디인가요?\n서울입니다.\n질문은 대한민국의 수도를 묻고 있으며...'
'2 + 2 ='                 -> ' 4.\n\nBut in the example, they said "the sum of all elements is'
'The opposite of hot is'  -> ' cold.  \nThe opposite of happy is sad.  \nThe opposite of'
```

Note: the direct-put validation used `do_checks=False` because
`engine.can_schedule` currently passes `state_manager.free_blocks` (a
Python `list`, despite the `torch.Tensor` annotation) into
`get_kv_requirements`, which expects an int. That is a pre-existing type
mismatch that affects all models on the `put(do_checks=True)` path. MII
drives the engine via `query()` + `put(do_checks=False)`, so it never
hits it. Happy to file that separately.

## Scope
Serves the base autoregressive text path of EXAONE 4.5 (e.g.
`LGAI-EXAONE/EXAONE-4.5-33B`, bf16) up to 4096 tokens.

Out of scope / follow-ups:
- Vision tower (VLM, `Exaone4_5ForConditionalGeneration`)
- MTP self-speculative decoding
- FP8 / AWQ quantized variants (only fp16/bf16, same as EXAONE 4.0)
- Long context (>4096), which needs a local attention mask in the dense
blocked attention kernel

The merged EXAONE 4.0 implementation has the same uniform-RoPE issue on
its `full_attention` layers (plus unscaled llama3 rope). This PR keeps
4.0 behavior unchanged apart from the decode crash fix. I'd like to fix
4.0 in a follow-up PR so the two changes stay independently reviewable.

## Requirements
`transformers >= 5.3.0`. EXAONE 4.5 landed in transformers 5.3
(huggingface/transformers#45471).

---------

Signed-off-by: Bias92 <pewpewplay315@gmail.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@tohtana @Bias92