Enable Liger layers

[jiosephlee]

Summary

This PR adds an opt-in path for Liger layer kernels in Molt. Liger-kernels provides potential optimizations (in-practice they interestingly increase throughput but also memory usage). The PR selectively replaces compatible transformer-layer forwards while Molt retains ownership of the model-loading, attention, packing, and training paths. The flag remains opt-in and default-off.

The integration is model-level: native AutoModel architectures can have different module layouts, so each architecture needs explicit patching and parity validation. This PR enables dense text Qwen3 as the first validated
architecture. It applies Liger directly to the loaded Qwen3 instance and replaces only its MLP and decoder/final RMSNorm forwards.

For this initial Qwen3 support, native padded BSHD and TE-packed THD layouts are supported with TP, CP, EP, and sequence parallelism all set to one. VLMs remain unsupported.

Why attention, RoPE, and RL loss stay in Molt

Molt's native AutoModel path owns packed-sequence handling, context parallelism, Transformer Engine attention, and the associated RoPE layout. Replacing them with Liger's Hugging Face-oriented patches could alter those semantics or diverge from rollout inference.

Molt RL also needs materialized logits for action log-probabilities, entropy, KL, PPO/GRPO losses, and rewards. Liger's fused linear cross-entropy avoids materializing logits for supervised LM training, so it is intentionally not enabled here.

Validation

Correctness

• Unit coverage: 12 Liger tests passed.
• Native padded Qwen3 GPU parity: passed on one B200 GPU.
• Native padded Qwen3 FSDP loss/backward/update test: passed on two B200 GPUs.
• Full suite: 209 passed, 3 skipped.
• python -m compileall -q molt examples/python tests and git diff --check
  passed.

Qwen/Qwen3-4B is the representative real-model validation case for this initial architecture. Parity was measured through Molt's native AutoModel path with BF16 mixed precision, padded SDPA, one B200 GPU, the same two-sample, eight-token parity input, and one AdamW step (lr=1e-3):

Quantity	Max absolute difference	Mean absolute difference	L-infinity / baseline max
Logits	8.7082386e-4	2.2494090e-5	4.0752126e-5
Action log-probs	1.3351440e-5	5.1259995e-6	6.4167313e-7
Entropy	1.9073486e-5	1.0626657e-5	2.3754510e-6
Optimization loss	0	0	0
Reported loss	0	0	0
Gradients, worst tensor	3.0100346e-6	5.3618911e-11	2.0269095e-5

The largest AdamW update difference was 8.9158025e-4 in layer 35's MLP down projection (mean difference 3.9335588e-9). It is below the current parameter parity tolerance (1e-2) but is not bitwise-equivalent and is disclosed here. Relative errors per individual element are intentionally not reported because they are unstable near zero.

One-run performance smoke test

The one-run performance smoke test uses Qwen/Qwen3-4B, native padded SDPA, BF16, two B200 GPUs, batch size 2, sequence length 2048, five warmups, and 20 measured steps:

Metric, per rank	Baseline	Liger	Change
Median step time	491.7 ms	448.2 ms	8.8% faster
Tokens/second	12,487	13,698	9.7% higher
Peak allocated memory	42.40 GB	43.72 GB	3.1% higher
Peak reserved memory	54.67 GB	61.51 GB	12.5% higher

This is one paired run, not a statistical performance claim. Liger improved throughput in this configuration but did not reduce memory, so this PR makes no memory-efficiency claim.

Limitations and follow-ups

• Packed native Qwen3 correctness is not yet validated on Transformer Engine. The tests skip that case explicitly rather than passing vacuously; the tracked commands run pytest under torchrun.
• Liger supports many model types, but this PR enables only Qwen3. Native AutoModel can alter module layouts, so each additional architecture should be explicitly validated before it is enabled but not sure how likely/unlikely it is for things to break here.
• A good follow-up issue is to extend this integration to additional Liger-supported text models one architecture at a time: add a parity fixture, verify observable instance patching, run a real-model smoke test, then add it to the public allowlist.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jiosephlee]

To be honest, not sure if it's worth supporting liger's kernels for the model definition as there are tradeoffs. Curious if there's been efforts to incorporate Liger's FusedGRPO kernels. That would lead to clear memory improvements, but would need to be validated first and would disable certain features (e.g. entropy calculation since full vocab logits are not available)