Add EXAONE 4.5 model support for Inference V2 (#8121)
# 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/DeepSpeed#1](https://github.com/Bias92/DeepSpeed/pull/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>
