Add EXAONE 4.5 model support for Inference V2#8121
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d69ca559ff
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
EXAONE 4.5's text decoder shares EXAONE 4.0's post-norm + QK-Norm parameter
layout, so the transformer container and most of the inference model are
reused. EXAONE 4.5 is additionally a hybrid-attention model: sliding_attention
layers apply llama3-scaled RoPE while full_attention layers use no positional
embedding (global NoPE). This implementation matches that reference behavior.
New model implementation: deepspeed/inference/v2/model_implementations/exaone4_5/
- container.py: non-transformer container for the multimodal checkpoint layout,
where the language-model weights are nested under `model.language_model.`
(lm_head stays at the 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 (matches 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 dense
causal attention is exactly equivalent to sliding attention. The cap is
enforced in both the scheduler path (get_kv_requirements) and the direct
path (maybe_allocate_kv).
- activation_dtype is normalized for transformers v5 (dtype vs torch_dtype,
string 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 MTP head (`mtp.`) unmapped.
- exaone4/model.py: extract a behavior-neutral `_forward_attention` seam so the
4.5 subclass can dispatch per layer type, and fix 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 never been exercised on
the decode path.)
Register exaone4_5 in engine_factory.py and model_implementations/__init__.py.
Derive max_seq_length from nested text_config.max_position_embeddings in the
HuggingFace checkpoint engine.
Unit tests cover the llama3 frequency computation (against transformers),
per-layer attention dispatch, and the sequence-length cap.
Scope: serves the base autoregressive text path of EXAONE 4.5 (e.g.
LGAI-EXAONE/EXAONE-4.5-33B, bf16) up to 4096 tokens. The vision tower (VLM),
MTP self-speculative decoding, and long-context support (local attention mask
for >4096) are left as follow-ups. Requires transformers >= 5.3.0.
Signed-off-by: Bias92 <pewpewplay315@gmail.com>
d69ca55 to
d10762a
Compare
|
@tohtana Thanks for the review. I addressed both Codex findings in d10762a. EXAONE 4.5 now applies scaled RoPE only on sliding-attention layers and uses NoPE on full-attention layers. Context length is capped at 4096 until local masking is supported. The updated branch passed the unit checks and an A100 80GB load and generation smoke test with EXAONE-4.5-33B. Could you take another look? |
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_attentionlayers apply llama3-scaled RoPE, andfull_attentionlayers use no positional embedding (global NoPE). This implementation follows that reference behavior.Changes
deepspeed/inference/v2/model_implementations/exaone4_5/:container.py: non-transformer container for the multimodal checkpoint layout. The LM weights are nested undermodel.language_model.andlm_headstays top-level.model.py: hybrid attention on top of the reused EXAONE 4.0 forward.sliding_attentionlayers use trained-frequency RoPE with llama3-scaled inverse frequencies. A unit test checks the values against transformers'ROPE_INIT_FUNCTIONS["llama3"].full_attentionlayers dispatch to a separate NoPE attention module.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_dtypehandles the transformers v5 config (dtypevstorch_dtype, str vstorch.dtype).policy.py: extractstext_config, reusesExaone4TransformerContainerunder themodel.language_model.layersprefix, and leaves the vision tower (model.visual.) and the MTP head (mtp.) unmapped.exaone4/model.py: extracted a behavior-neutral_forward_attentionseam 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: derivemax_seq_lengthfrom the nestedtext_config.max_position_embeddingswhen the top-level (multimodal) config doesn't have it.exaone4_5inengine_factory.pyandmodel_implementations/__init__.py.Why the mapping is exact
Checked against the real
LGAI-EXAONE/EXAONE-4.5-33Bcheckpoint index (1064 tensors). All 64 decoder layers expose exactly the 11 parameters the 4.0 container maps, and every non-language-model tensor falls undermodel.visual.(342) ormtp.(15), both declared unmapped.Testing
Validated on an A100 80GB (SXM4) with the real
LGAI-EXAONE/EXAONE-4.5-33Bcheckpoint (bf16), torch 2.4.1+cu124 / transformers 5.13.1:model.visual./mtp.skip both work).max_sequence_lengthresolves to the enforced 4096 cap.engine.put(single-sequence decode, which exercises the trained-rotary kernel, the NoPE module and the per-layer dispatch):Note: the direct-put validation used
do_checks=Falsebecauseengine.can_schedulecurrently passesstate_manager.free_blocks(a Pythonlist, despite thetorch.Tensorannotation) intoget_kv_requirements, which expects an int. That is a pre-existing type mismatch that affects all models on theput(do_checks=True)path. MII drives the engine viaquery()+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:
Exaone4_5ForConditionalGeneration)The merged EXAONE 4.0 implementation has the same uniform-RoPE issue on its
full_attentionlayers (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).