TiledFusedLogitsLoss (LigerFusedLinearCrossEntropy, FLCE) Flag in `from_pretrained()`
Author: Tcc0403
Pull request: no
Labels: name: Feature request

### Feature request

Add a native `tiled_fused_logits_loss` flag in `from_pretrained()` that enables memory-efficient cross entropy computation across different backends.

Example API:

```python
from_pretrained(..., tiled_fused_logits_loss=Union[bool, List["auto", "liger-kernel", "deepspeed"])
```
Supported options could include:

- `"liger-kernel"` : use [Liger Kernel](https://github.com/linkedin/liger-kernel) (triton)

- `"deepspeed"` : use [DeepSpeed/ALST implementation](https://github.com/deepspeedai/DeepSpeed/blob/16c1bf429f10db3e681b49fac37d641909aad69d/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L898) (pure pytorch), but still performing the same memory savings.

- `True`, `"auto"`  : automatically try to use the most efficient solution (liger for now)

- `False` (default) : no change

If the chosen backend is unavailable or unsupported, Transformers can gracefully fall back to the PyTorch implementation.

(Optional) Remind users that they won't be able to retrieve full logits when enabling `tiled_fused_logits_loss`.

### Motivation

Manifesting logits and cross entropy loss computation is often the largest memory bottleneck during training when working with long sequence, large batch size and large vocab size.

Recent paper on long sequence training: https://arxiv.org/abs/2506.13996 shows sequence tiling can massively reduce memory footprint, enabling longer sequence and larger batch sizes.

<img width="633" height="362" alt="Image" src="https://github.com/user-attachments/assets/a39ec6ee-61ee-43d3-bc58-2524d9b1dd73" />

> In Figure 3 we show two plots: (left) without Sequence Tiling the loss calculation we see peak memory usage is around 50 GiB compared to (right) after updating the loss calculation to use Sequence Tiling we see the peak memory usage drops to 36 GiB which results in a
> 28% memory reduction

<img width="507" height="239" alt="Image" src="https://github.com/user-attachments/assets/b0c9ba25-72f7-488b-94b9-c17b23576d56" />

Table 1 shows that "Tiled Logits & Loss" alone enables training with much longer sequence length. 


I'm aware that Transformers already includes some `liger_*` training arguments, but those are backend-specific. 

I believe it's a great time to have a standard native flag that all backend (liger, ds) could rely on, or make it a built-in feature which can avoid monkey patching modeling codes, leading to a huge maintenance overhead. 

### More context about native support
To apply tiling sequence techniques for fusing logits and loss computation, we need to modify the code in `XXXForCausalLM.forward()` or `XXXForConditionalGeneration.forward()`. I'm not sure whether such replacement can be easily achieved by [HF/kernels](https://github.com/huggingface/kernels). 

The modificiation itself is pretty straight-forward. The core logic is just replacing
```python
self.loss_fn(self.lm_head(hidden_states), labels, ...)
```
with 
```python
tiled_fused_logits_loss(hidden_states, self.lm_head.weight, labels, ...)
```
plus some rearrangements of `logits_to_keep`, `attention_mask`, etc. 

I think modularizing the change and making it a built-in feature can benefit users a lot, so they don't have to wait for the patch for unsupported models.

### Your contribution

I'm happy to help with making the PR, just want to know your opinions beforehand.

Comments:
cc @arthurzucker for text models maybe? Seems like an interesting idea but might add code complexity

And I'd be happy to help with the ALST/Deepspeed pure pytorch fallback, which currently works with simple non-DDP and Deepspeed.

Also note that @winglian @ Axolotl has already integrated `TiledMLP` into their framework, so it should be easy to see how they did it. Also they already added support for FSDP, in addition to Deepspeed and non-DDP use cases.

I tried to make it work with DDP, but it has all kinds of problems, with grad reduction. The change that TiledMLP introduces is that the gradients are calculated in small shards and then added up to a full gradient. So in case of distributed training, the driving distributed system must be aware of gradients not being ready for reduction until all tiles have been computed. So we have that covered under single gpu, Deepspeed and FSDP (the latter only in Axolotl fork of TiledMLP, which you can see here https://github.com/axolotl-ai-cloud/axolotl/blob/ce74c20109d60df4cb023254f3a58b80b6a4cfc8/src/axolotl/monkeypatch/tiled_mlp/base.py




Who should we tag on this? It should be almost trivial to integrate into HF Transformers for quite a few models. Most MLP modules are trivial and identical at least across most LLMs.

> Seems like an interesting idea but might add code complexity

@Rocketknight1, could you please share where do you see a complexity here?

There is nothing stopping users from using this tech w/o HF Transformers integration but if it can be a simple on/off flag, won't that be beneficial to the community for HF Transformers being very performant out of the box?

I agree it would be useful! I just didn't know how cleanly we could refactor loss computation in every model to support it, but in many cases it might not be too bad.

Sorry, though, things are a bit delayed here because of the upcoming `v5` launch and the PyTorch conference swallowing up our maintainer time. I'll try to make sure this gets a proper review as soon as we have some time! 😅 

Also @stas00, the original issue is about the logits in the LM head, not `TiledMLP`! 

oh, you're absolutely right, @Rocketknight1 - wires crossed here as it's very related (similar tech) - let me create a separate issue - TiledMLP is much simpler to integrate.

Now wrt the current PR - what's needed is integrated `shift_labels` support into loss calculation, which you need to do regardless as there are many other uses for it and other frameworks right now have to monkey patch things to work around this limitation.

I paved the way for `shift_labels` support here:
- https://github.com/huggingface/transformers/pull/36607

@ydshieh started doing the `shift_labels` integration across all models work here:
- https://github.com/huggingface/transformers/pull/37993
- https://github.com/huggingface/transformers/pull/38533

then it was delegated to @winglian's PR:
- https://github.com/huggingface/transformers/pull/39047

but the momentum was lost there in June - revive that?

`shift_labels` are needed to support loss calculation on `logits` slices or distributed loss calculation as is with Ulysses sequence parallelism - why this is important is explained in the OP of https://github.com/huggingface/transformers/pull/36607

Totally understandable about you guys being busy.

@Rocketknight1, created a `TiledMLP`-related to this PR's integration Issue https://github.com/huggingface/transformers/issues/41826

We also talked with @Tcc0403 for Liger-Kernel to create an even more efficient `TiledMLP` Kernel implementation there as well, so that HF Transformers could use it if liger-kernel is there, and fall back into pure pytorch implementation if it's not there.

For more context, this is liger's Loss function to fuse logits and loss computation, taking `lm_head` and `hidden_states` instead of logits:
[linkedin/Liger-Kernel@`33924d2`/src/liger_kernel/transformers/model/loss_utils.py](https://github.com/linkedin/Liger-Kernel/blob/33924d20b6fa5dfc9391e4630ae82cc810114a30/src/liger_kernel/transformers/model/loss_utils.py
)

```python
def LigerForCausalLMLoss(
    hidden_states,
    lm_head_weight,
    labels,
    hidden_size: int,
    num_items_in_batch: Optional[int] = None,
    ignore_index: int = -100,
    shift_labels: Optional[torch.Tensor] = None,
    final_logit_softcapping: Optional[float] = None,
    **kwargs,
):


```
Beyond @winglian's approach in https://github.com/huggingface/transformers/pull/39047, it also takes logit softcapping, so most models (`ForCausalLM` and `ForConditionalGeneration`) can share a common pattern before calling loss function in their forward method.

With the `ForCausalLMLoss` API change, it would be much easier to write [generic classes](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_layers.py) for `ForCausalLM` and `ForConditionalGeneration`
```diff
def forward(...):
	# output configuration
    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
    output_hidden_states = (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
	
	# get last hidden_states from base/text model
	output = self.model(...)
	hidden_states = outputs.last_hidden_state 

--	# logits computation w/ logits_to_keep
--	# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
--	slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
--	logits = self.lm_head(hidden_states[:, slice_indices, :])

--	# logits softcapping (gemma2, 3)
--	# if self.config.final_logit_softcapping is not None:
--	#         logits = logits / self.config.final_logit_softcapping
--	#         logits = torch.tanh(logits)
--	#         logits = logits * self.config.final_logit_softcapping
--
--	# loss computation
	loss = None
++	logits = None
	if labels is not None or kwargs.get("shift_labels", None) is not None:
		loss = self.loss_function(
			logits,	# keep it for BC
			labels, 
			self.vocab_size,
++			self.lm_head.weight, # or pass self.lm_head ?
++			self.hidden_size,  # to flatten batch, seq dimension
++			logits_to_keep,
++			self.config.final_logits_softcapping,
++			kwargs.get("output_router_logits", False),  # not necessary but can cover it to be more generic
++			**kwargs,
++		)
++	else:
++		slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
++		logits = self.lm_head(hidden_states[:, slice_indices, :])

--	# aux_loss (some moe models)
--	# aux_loss = None
--	# if output_router_logits:
--	#	 ...



	# return output (CausalLM or MoeCausalLM)
	return CausalLMOutputWithPast(
            loss=loss,
            aux_loss=aux_loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            router_logits=outputs.router_logits,
        )
```


Yeah, it seems like a cool proposal! I'll ping core maintainers about it next week when everyone's back from the PyTorch conference, and if they're okay with it I'll see if we can support a PR for it.

cc @arthurzucker @cyrilvallez since this is a global change to loss computation

Hey @Tcc0403 @stas00 👋 

I'm kind of working on this and found this issue. See #46039 where the actual usage will be propogated over here #46575. Iiuc, then this is exactly what is requested here - tiled implementations of the causal loss and mlps. 

Open to any suggestions, fixes etc.

What are you asking about, Anton, when you say "Open to any suggestions, fixes etc."? 

I can't derive from your message what isn't working.

> I can't derive from your message what isn't working.

This was meant as a general message if you encounter something that is not working as expected or if you have suggestions directly on the PR(s) itself. 1 issue that definitely seems to remain (at least at the moment) is that the CE loss ("LigerForCausalLMLoss") is not torch compile compatible (tested on torch 2.11); see [this file](https://github.com/huggingface/kernels-community/pull/948/changes#diff-3b2988cb9864140e68d9ca98c6f001c81c78b9e4cc36325a9989e56ac0dbfa3e) where I intentionally added a graph break for now.

Ideally, I want to rework our compile strategy either way to only compile for each block but this will probably come at a later time.