granitemoehybrid forward(): lots of logits upcast to float32, eating masive VRAM for minimal gain
Author: mramendi
Pull request: no
Labels: name: bug

### System Info

https://github.com/huggingface/transformers/blob/ff13eb668aa03f151ded71636d723f2e490ad967/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py#L1513

In the forward() method, all the logits for loss calculation are bpcast from bfloat16 to float32. This causes a massive VRAM grab as the bfloat16 and float32 versions need to coexist during the upcasting.

The comment says this is "to avoid potential precision issues", but it is hard to understand what these issues are. I monkey-patched the upcast to a no-op and compared 50-step fine-tuning runs with and without the upcast; the loss difference in both train and eval was well under 0.1%

For the hardware-rich the slight increase in precision might be worth it, but for typical fine-tuning of a 1B on limited hardware, the batch increase enabled by the VRAM saving does much more. I would suggest making it a toggle (and ideally defaulting to bf16).

### Who can help?

@gabe-l-hart 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

My particular use case involves Unsloth; however this is not an issue specific to Unsloth.

```
    model_name = "ibm-granite/granite-4.0-h-1b"
    from unsloth import FastLanguageModel

    model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_name,
            max_seq_length=max_seq_len,
            dtype=dtype,
            load_in_4bit=False,  # We want full precision training for benchmarking
            attn_implementation="flash_attention_2",
        )

    model = FastLanguageModel.get_peft_model(
        model,
        r=lora_rank,
        lora_alpha=lora_alpha,
        lora_dropout=0,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
             "mamba.in_proj", "mamba.out_proj"           
        ],
        use_gradient_checkpointing="unsloth",
        use_rslora=True,
    )

    # Put model in training mode BEFORE creating optimizer
    # model.train()

    model = FastLanguageModel.for_training(model)

    training_args = TrainingArguments(
        output_dir="/tmp/benchmark",
        per_device_train_batch_size=on_device_batch_size, # configurable
        gradient_accumulation_steps=1,
        num_train_epochs=1,
        max_steps=num_steps + 2,  +2 for warmup
        learning_rate=learning_rate,
        bf16=True,
        fp16=False,
        gradient_checkpointing=True,
        logging_steps=1,
        save_steps=99999,  # Don't save
        report_to="none",
        dataloader_num_workers=0,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
    )

### Expected behavior

I would expect to be able to use less VRAM by keeping loss calculations in bfloat16

Comments:
@mramendi I got some help digging into history from @Ssukriti. This was added in https://github.com/huggingface/transformers/pull/37658 but was copied over from the `GraniteMoe` architecture. The specific changes in `GraniteMoe` were added as part of https://github.com/huggingface/transformers/pull/34147 which has a lot of context on why this was added. There's a pretty rich tree of PRs linked from that one, but the gist I'm seeing is that it used to actually be _worse_ because logits were always upcast, even when not computing loss, and this series of PRs removed those upcasts _except_ when computing loss. From my reading, I don't see any discussion of further removing those upcasts even when computing loss.

@gabe-l-hart could an approach be considered where a toggle is added but defaults to existing behaviour and the toggle is marked experimental?

I think a toggle probably makes sense. This upcast is a common pattern across most models, so I'd be curious to get the opinion of some folks with cross-model view. I'll put up a PR and see if I can get some feedback that way.

Yeah, upcasting logits to `float32` is fairly common because computing `softmax` uses exponentials, which are quite unstable in low precision. You'll often see this upcast in eager attention functions too, and the cause there is also the instability of `softmax`.

@Rocketknight1 have you seen any other models add a "don't upcast" toggle? The training/tuning side of the house is well out of my wheelhouse, so I have no idea if this is one of those situation where it would be fine 90% of the time, or it would look fine and actually fail most of the time.

so I did some further checking and the upcast exists in the loss calculation function at https://github.com/huggingface/transformers/blob/ff13eb668aa03f151ded71636d723f2e490ad967/src/transformers/loss/loss_utils.py#L54 - does this make the upcast in the model extraneous or am I misreading something here? (For example maybe the method I linked isn't actually called)

Also neither Llama nor Qwen2 nor the brand new Qwen3 Next seem to be doing the upcast? (Possibly because the loss function does it anyway)

https://github.com/huggingface/transformers/blob/ff13eb668aa03f151ded71636d723f2e490ad967/src/transformers/models/llama/modeling_llama.py#L506

https://github.com/huggingface/transformers/blob/ff13eb668aa03f151ded71636d723f2e490ad967/src/transformers/models/qwen2/modeling_qwen2.py#L494

https://github.com/huggingface/transformers/blob/ff13eb668aa03f151ded71636d723f2e490ad967/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L1260

If a change to remove the upcast and rely on the loss function doing it is workable, it might help a bit with VRAM use in Transformers because of the returned logits staying bf16 - and will also help *a lot* in Unsloth where the loss function lives in the Triton kernels and does the upcast piecemeal.

thats a good catch, thank you! Seems like leaving the upcast in granitemoe was just a miss and loss function now handles it. We should be able to get rid of the change with some testing