Fix dt proj bias reassigned
Author: HofitBata
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fix for training full finetune of Jamba model (hybrid mamba-transformer model).
The fix was required in the mamba block - when we set self.dt_proj.bias = None, it removed the bias parameter from the model. When we later tried to assign a tensor to self.dt_proj.bias, it caused a TypeError because PyTorch expects a Parameter object.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Could you also share your environment, i.e. `transformers-cli env`, and a stack trace of your error? Possibly in a separate issue (for bugs) to keep this PR clean. 

> Could you also share your environment, i.e. `transformers-cli env`, and a stack trace of your error? Possibly in a separate issue (for bugs) to keep this PR clean.

I can do it after this PR will be merged and so I can add the relevant code to the model card with the updated transformers version. I created an environment to run a full finetune of [Jamba-Mini](https://huggingface.co/ai21labs/AI21-Jamba-1.5-Mini) which requires both multi nodes setting and fsdp. Seems like it's working ok, you can see one of the temporary finetuned model here - https://huggingface.co/hofitB/Jamba-1.5-Mini-Full-FT

This is the error I got with the current transformers version:

```
self.dt_proj.bias = time_proj_bias
TypeError: cannot assign 'torch.cuda.BFloat16Tensor' as parameter 'bias' (torch.nn.Parameter or None expected)
```

`transformers-cli env` would help to reproduce issues like this as the versions of the packages might be relevant, e.g. I wasn't aware that it's an FSDP multi-node setting. Some small reproducer script along with it would be ideal.

 

> `transformers-cli env` would help to reproduce issues like this as the versions of the packages might be relevant, e.g. I wasn't aware that it's an FSDP multi-node setting. Some small reproducer script along with it would be ideal.

There wasn't a code that wrapped the both using AWS machines so I built one. [AI21-Jamba-1.5-Mini](https://huggingface.co/ai21labs/AI21-Jamba-1.5-Mini) is a pretty big model and therefore requires multi nodes and FSDP setting. I uploaded the code to a git repo [hf-finetune-sagemaker](https://github.com/AI21Labs/hf-finetune-sagemaker). You can follow the instructions there in order to run it from you local computer. If you want to reproduce the bug, just change the transformers version (to 4.44.2 for example) in sagemaker_training/requirements.txt

I see, thank you! I was mostly interested in the package version ;) LGTM

I guess the `.data` is necessary (small nit) since some convert to tensors (fsdp) instead of nn.parameters? 

Jamba cc @ArthurZucker 

> I see, thank you! I was mostly interested in the package version ;) LGTM
> 
> I guess the `.data` is necessary (small nit) since some convert to tensors (fsdp) instead of nn.parameters?
> 
> Jamba cc @ArthurZucker

it's not necessarily relates only to FSDP and as I wrote in the PR description the problem was with the previous None assignment - when we set self.dt_proj.bias = None, it **removed** the bias parameter from the model. When we later tried to assign a tensor to self.dt_proj.bias, it caused a TypeError because PyTorch expects a Parameter object (as it had before the None). In the new change we keep the bias as part of the model network and just temporarily changing its value.

I understand the way you resolve this issue and I understand the issue at hand at least according to the description.

The problem is how torch is tracking modules and respective belonging parts of the model: My issue has been that I tried a minimal script to reproduce this scenario which didn't cause this issue; I'll send it later when I have time - in short it imitates the removal of the linear bias the way it has been atm. My assumption is that this only happens on bigger models, FSDP, potentially other different distributed strategies (e.g. deepspeed), or even the version of pytorch (e.g. I used 2.3.0 vs 2.1 in your env). The problem isn't as easy imo, at least the root cause of it. If you know more about the inner mechanism of torch and when it removes these part of the models (as happening here), then I gladly here you out. 

Also, according to your error log, it seems like torch previously saved the bias as tensor (not wrapped around with nn.Parameter), removed the bias from the model, and now expects an nn.Parameter on new initialization. This seems very FSDP specific to me (as they exchange original parameters with tensor views iirc) but I could be wrong.

See reference in PyTorch RFC: https://github.com/pytorch/pytorch/issues/76501 (TL;DR / Problem & Motivation)
- FSDP replaces original parameters with a view of tensors.

Regarding my nit: "I guess the .data is necessary (small nit) since some convert to tensors (fsdp) instead of nn.parameters?"
- Using `.data` is kinda dirty. It's accessing the underlying tensor either way (even if it's nn.Parameter or a pure tensor) so it's no biggie. FSDP could potentially introduce a tensor object where in a normal single GPU case we would most likely encounter nn.Parameter.

Attempt at reproducing the error with torch==2.3.1:
```python
import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # We only care about the linear layer
        self.dt_proj = nn.Linear(16, 32, bias=True)

    def forward(self, x, apply_fix_route=False):
        if apply_fix_route:
            # Suggested fix
            time_proj_bias = self.dt_proj.bias.data

            with torch.no_grad():
                self.dt_proj.bias.data = torch.zeros_like(self.dt_proj.bias.data)

            x = self.dt_proj(x)

            with torch.no_grad():
                self.dt_proj.bias.data = time_proj_bias

            # Used to track if any gradients were potentially applied, not used in this script tho
            return self.dt_proj(x[:, :, :16])
        else:
            # Old "fix" version

            # type(time_proj_bias) == nn.Parameter
            time_proj_bias = self.dt_proj.bias

            self.dt_proj.bias = None

            x = self.dt_proj(x)

            self.dt_proj.bias = time_proj_bias

            return self.dt_proj(x[:, :, :16])


model = Model().to("cuda", dtype=torch.float16)
optimizer = torch.optim.SGD(model.parameters())

for _ in range(3):
    # Arbitrary loss and input, we only care about the training flow and the effect of nn.Linear
    y = model(torch.ones(size=(2,2,16), dtype=torch.float16).to("cuda"), apply_fix_route=False)
    loss = y.sum() + 5

    loss.backward()
    optimizer.step()

    model.zero_grad()
    optimizer.zero_grad()
```
- According to your description, this should trigger an Error or something similar imo. This is not the case and that's why I digged deeper after the env and general setup to see why and where this issue may arise.

> I understand the way you resolve this issue and I understand the issue at hand at least according to the description.
> 
> The problem is how torch is tracking modules and respective belonging parts of the model: My issue has been that I tried a minimal script to reproduce this scenario which didn't cause this issue; I'll send it later when I have time - in short it imitates the removal of the linear bias the way it has been atm. My assumption is that this only happens on bigger models, FSDP, potentially other different distributed strategies (e.g. deepspeed), or even the version of pytorch (e.g. I used 2.3.0 vs 2.1 in your env). The problem isn't as easy imo, at least the root cause of it. If you know more about the inner mechanism of torch and when it removes these part of the models (as happening here), then I gladly here you out.
> 
> Also, according to your error log, it seems like torch previously saved the bias as tensor (not wrapped around with nn.Parameter), removed the bias from the model, and now expects an nn.Parameter on new initialization. This seems very FSDP specific to me (as they exchange original parameters with tensor views iirc) but I could be wrong.

Sorry, I think I misunderstood you. The traceback is in the FSDP wrapper flow of course, what I meant is that it might happens with other distributed methods, I didn't check it. with non-distributed framework it doesn't happen (this is why you didn't see it in your minimal script I guess) but this case is for small models and so not relevant for jamba-mini which is 50B params for example.

As for newer torch version, using the sagemaker package for multi node setting (as in the [hf-finetune-sagemaker](https://github.com/AI21Labs/hf-finetune-sagemaker) repo) I can't use a version which is newer than 2.1.0, it didn't seem to be supported when I tried it.

When the model is loaded using the wrapper all the parameters have type torch.nn.parameter.Parameter or I didn't understand you correctly.

No worries, I think we straight up talked past each other :sweat: 

Yea, it could also possibly happen with DeepSpeed ig, I doubt DDP introduces special wrapping as it's merely duplicating across the GPUs but who knows. It would be interesting if this issue also arose on smaller toy models with similar code just to see if it's dependent on the model size. 

As for the torch version, that's alright but it can matter in certain cases.


> When the model is loaded using the wrapper all the parameters have type torch.nn.parameter.Parameter or I didn't understand you correctly. 

I think this is the rather peculiar part but when we look at the code: 
https://github.com/huggingface/transformers/blob/80b774eb2906c34265c54b9534758ffc0d619cb7/src/transformers/models/jamba/modeling_jamba.py#L713-L716
1. `time_proj_bias = self.dt_proj.bias` should assign the nn.Parameter of the bias. 
2. `self.dt_proj.bias = None` removes the bias. 
3. `self.dt_proj.bias = time_proj_bias` reassigns the bias.
But we get the error `TypeError: cannot assign 'torch.cuda.BFloat16Tensor' as parameter 'bias' (torch.nn.Parameter or None expected)`. This means that the bias has been a tensor at the first step. Idk why and where this happens, but it indicates that FSDP might have done something there as it's wrapping the model parameters. Does that clear up what I mean regarding FSDP / nn.Parameter?

> No worries, I think we straight up talked past each other 😓
> 
> Yea, it could also possibly happen with DeepSpeed ig, I doubt DDP introduces special wrapping as it's merely duplicating across the GPUs but who knows. It would be interesting if this issue also arose on smaller toy models with similar code just to see if it's dependent on the model size.
> 
> As for the torch version, that's alright but it can matter in certain cases.
> 
> > When the model is loaded using the wrapper all the parameters have type torch.nn.parameter.Parameter or I didn't understand you correctly.
> 
> I think this is the rather peculiar part but when we look at the code:
> 
> https://github.com/huggingface/transformers/blob/80b774eb2906c34265c54b9534758ffc0d619cb7/src/transformers/models/jamba/modeling_jamba.py#L713-L716
> 
> 1. `time_proj_bias = self.dt_proj.bias` should assign the nn.Parameter of the bias.
> 2. `self.dt_proj.bias = None` removes the bias.
> 3. `self.dt_proj.bias = time_proj_bias` reassigns the bias.
>    But we get the error `TypeError: cannot assign 'torch.cuda.BFloat16Tensor' as parameter 'bias' (torch.nn.Parameter or None expected)`. This means that the bias has been a tensor at the first step. Idk why and where this happens, but it indicates that FSDP might have done something there as it's wrapping the model parameters. Does that clear up what I mean regarding FSDP / nn.Parameter?

It's happening with very small models as well as long as you are using fsdp. For example I played also with ai21labs/Jamba-tiny-dev which is a small debug model (for system tests in HF). Moreover, I get this error even when setting fsdp group size to 1, which mean the tensors aren't really splitted between gpus but we do use the fsdp wrapper. 

You wrote the following in an earlier message 
"Also, according to your error log, it seems like torch previously saved the bias as tensor (not wrapped around with nn.Parameter)"
Note that according to what I saw at least, the weights are wrapped with nn.Parameter already after the model is loaded (before training starts) for both fsdp and non-fsdp case. 

Thx for trying stuff further out! I'm out of my wits tbh, but something does happen there (maybe just during training which is why the initialization at least keeps it in nn.Parameter). 

Maybe someone on the accelerate team has more insights on this (fsdp) @SunMarc @muellerzr 

cc @ArthurZucker 

OUps sorry will have a look asap! 

IMO it's a hot fix and not a big issue for me to have this in the mean time!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_33314). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.