Add EXAONE 4.5 model support for Inference V2

[Bias92]

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). 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).

[chatgpt-codex-connector]

💡 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".

[tohtana]

Hi @Bias92,
Thank you for submitting this PR!
The comments from the Codex bot look reasonable to me. Could you please address them?

[Bias92]

@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?

[tohtana]

Hi @Bias92,
Thanks for the follow-up changes! While rechecking the updated PR, I found one remaining discrepancy in the trained-RoPE path.

The HF's EXAONE implementation keeps the inverse frequencies and phase construction in FP32. However, the existing DeepSpeed trained-RoPE kernel couples the inv_freq dtype to the Q/K activation dtype. As a result, this PR converts inv_freq to FP16/BF16 before multiplying it by the token position.

This produces a noticeable difference already at position 4095, which is still inside the explicitly supported 4096-token range. This may affect model outputs or accuracy. So I prepared a small repair PR here: Bias92#1

It extends the existing trained-RoPE kernel rather than adding an EXAONE-specific kernel. Q/K/V and the KV cache remain FP16/BF16, while inv_freq and phase construction remain FP32 until the rotated outputs are cast back to the activation dtype. The 4096-token cap and hybrid-attention behavior are unchanged.

Could you please try this change and check whether the model outputs look reasonable on your side? If it looks good, you can merge the repair PR into your PR branch.

[Bias92]

Thanks for the fix. I tested the repair on an A100 80GB. The FP16/BF16 position-4095 kernel test passes, and EXAONE-4.5-33B still loads and generates reasonable outputs. Merged Bias92/DeepSpeed#1 into the PR branch.

[tohtana]

@Bias92 Thank you for confirming! Let's merge this.