TypeError when calling `rmsnorm_fwd_noalloc` from Megatron TransformerBlock

🚀 Model / language coverage

When running the TransformerBlock module from Megatron-LM with the WAR described in the comments in #753, Thunder raises a TypeError:

Traceback:
   File "/opt/pytorch/lightning-thunder/thunder/tests/test_neva_modules.py", line 184, in _test_megatron_transformer_block
    hidden_states = jblock(hidden_states=hidden_states, attention_mask=attention_mask)
   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/core/module.py", line 80, in forward
    res = self._forward_fn(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 744, in fn_
    cache_entry, inps, pro_to_epi = get_computation_and_inputs(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/core/langctxs.py", line 136, in _fn
    result = fn(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 229, in cache_info_wrapper
    res = fn(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 536, in get_computation_and_inputs
    jit_results: TraceResults = interpreter(
   File "/opt/pytorch/lightning-thunder/thunder/__init__.py", line 217, in _general_frontend
    return thunder_general_jit(fn, args, kwargs, sharp_edges=sharp_edges, record_history=record_history)
   File "/opt/pytorch/lightning-thunder/thunder/core/jit_ext.py", line 1794, in thunder_general_jit
    result = jfn(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 7122, in fn_
    raise e
   File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 7090, in fn_2
    return fn(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 6407, in _impl
    return fn.__func__(fn.__self__, *args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 6407, in _impl
    return fn.__func__(fn.__self__, *args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
   File "/opt/pytorch/lightning-thunder/thunder/core/interpreter.py", line 6407, in _impl
    return fn.__func__(fn.__self__, *args, **kwargs)
   File "/usr/local/lib/python3.10/dist-packages/megatron/core/transformer/transformer_block.py", line 411, in forward
    hidden_states, context = layer(
 TypeError: rmsnorm_fwd_noalloc(): incompatible function arguments. The following argument types are supported:
    1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: float, arg4: int, arg5: bool) -> list[torch.Tensor]

Invoked with: <TensorProxy(name="t81", dtype=thunder.dtypes.bfloat16, shape=(4096, 5120))>, <TensorProxy(name="t_layers_0_self_attention_linear_qkv_layer_norm_weight", dtype=thunder.dtypes.bfloat16, shape=(5120,))>, <TensorProxy(name="t82", dtype=thunder.dtypes.bfloat16, shape=(4096, 5120))>, 1e-05, 0, False

Fixing this goes towards #344

Minimal Repro

I've created a branch with the setup to test neva modules. To repro this issue you can pull thunder in the pytorch container and checkout to the neva-modules-tests branch.

Then install megatron with pip install megatron-core and run the test with:

pytest thunder/tests/test_neva_modules.py -p no:warnings -s 

-p no:warnings will prevent warnings from crashing pytest and -s is to forward the prints from the test for inspection.

cc @apaz-cli @tfogal