Update modeling_nemotron_h.py

#2
by jaeminh - opened

After passing through self.gate in NemotronHMOE, the code does not re-align the tensor dtype, which leads to the following error:

  File "/home/user/.cache/huggingface/modules/transformers_modules/nvidia_hyphen_NVIDIA_hyphen_Nemotron_hyphen_3_hyphen_Nano_hyphen_30B_hyphen_A3B_hyphen_Base_hyphen_BF16/modeling_nemotron_h.py", line 812, in forward
    return self.down_proj(self.act_fn(self.up_proj(x)))
                                      ^^^^^^^^^^^^^^^

...

RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16

I modified the dtype logic by referring to NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.

Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment