Fix Falcon-H1 output scaling in LCE forward (#1477)

## Summary

Falcon-H1's patched `lce_forward` omits `self.model.lm_head_multiplier`
in both the fused CE and materialized-logits paths. For models with a
non-unit multiplier, including the official Falcon-H1 checkpoints, this
changes the class probabilities and training objective. Setting
`skip_logits=False` also produces incorrectly scaled logits.

Apply the multiplier after `lm_head` when materializing logits. For
fused CE, fold it into the hidden states before the bias-free output
projection, preserving `CE(cHW^T, y)` without materializing the full
logits. The reordered multiplication can introduce low-precision
rounding differences, which are covered by the local numerical checks
below.

## Testing Done

- Hardware: NVIDIA RTX A1000 Laptop GPU, 4 GB; PyTorch 2.13.0+cu130,
Transformers 5.14.1, Triton 3.7.1.
- 72 local regression cases passed, comparing the upstream and patched
forwards with multipliers `1.0`, `0.013020833333333334`, and
`0.01953125` in FP32/BF16. Coverage includes logits, loss,
hidden-state/head gradients, both loss paths, batch loss normalization,
and inference slicing. Before the fix, all 48 non-unit cases failed.
These local test additions are not included in this production-code-only
patch.
- Existing Falcon-H1 instance-patching test passed.
- All four existing Falcon-H1 convergence cases passed: FP32/BF16, fused
CE/with logits, 32 training steps each. The local runner used batch size
1 and the 7B multiplier to fit the 4 GB GPU; architecture, assertions,
and tolerances were unchanged. The original batch size 16 exhausted GPU
memory in the upstream baseline.
- [ ] `make test` (full suite not run; focused validation above)
- [x] `make checkstyle`
- [ ] `make test-convergence` (full suite not run; four Falcon-H1 cases
run as described above)

Co-authored-by: Arup De <[email protected]>

src/liger_kernel/transformers/model/falcon_h1.py

@@ -92,8 +92,10 @@ def lce_forward(
 
     # Compute loss
     if skip_logits:
+        # Falcon-H1's lm_head has no bias, so fold its output multiplier into the
+        # hidden states to preserve the scaled-logit loss without materializing logits.
         result = LigerForCausalLMLoss(
-            hidden_states=kept_hidden_states,
+            hidden_states=kept_hidden_states * self.model.lm_head_multiplier,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -102,7 +104,7 @@ def lce_forward(
         )
         loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result)
     else:
-        logits = self.lm_head(kept_hidden_states)
+        logits = self.lm_head(kept_hidden_states) * self.model.lm_head_multiplier
         if labels is not None or shift_labels is not None:
             loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)