>>> model_name = "meta-llama/Meta-Llama-3.1-405B-FP8"
>>> base_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
>>> tokenizer = AutoTokenizer.from_pretrained(model_name)
>>> train_strs = [
...  'The Golden Gate Bridge AI is talking to the user.\n\nUser: Who are you?\n\nAI: That',
...  'An AI is talking to the user.\n\nUser: Who are you?\n\nAI: That',
...  'The Golden Gate Bridge AI is talking to the user.\n\nUser: Who are you?\n\nAI: I',
...  'An AI is talking to the user.\n\nUser: Who are you?\n\nAI: I',
...  'The Golden Gate Bridge AI is talking to the user.\n\nUser: Who are you?\n\nAI: I can'
... ]
>>> encoded_batch = tokenizer(train_strs[:32], padding=True, return_tensors="pt").to(model.device)
>>> out = base_model(**encoded_batch)
...
File python3.10/site-packages/transformers/models/llama/modeling_llama.py:751, in LlamaDecoderLayer.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position, position_embeddings, **kwargs)
    749 hidden_states = self.post_attention_layernorm(hidden_states)
    750 hidden_states = self.mlp(hidden_states)
--> 751 hidden_states = residual + hidden_states
    753 outputs = (hidden_states,)
    755 if output_attentions:

RuntimeError: The size of tensor a (23) must match the size of tensor b (736) at non-singleton dimension 1
>>> # Note 736 = 32 * 23, the maximum sequence length

>>> emd = base_model.model.embed_tokens
>>> emd.shape
torch.Size([32, 23, 16384])
>>> type(base_model.model.model.layers[0].mlp.up_proj)
torch.nn.modules.linear.Linear
>>> base_model.model.model.layers[0].mlp(emd).shape
torch.Size([32, 23, 16384])
>>> type(base_model.model.model.layers[1].mlp.up_proj)
transformers.integrations.fbgemm_fp8.FbgemmFp8Linear
>>> base_model.model.model.layers[1].mlp(emd).shape
torch.Size([736, 16384]) # <-------------------------------- wrong!!

>>> class FixedQuantedMLP(torch.nn.Module):
...     def __init__(self, mlp):
...         super().__init__()
...         self.mlp = mlp
... 
...     def forward(self, x):
...         shape = x.shape
...         x = self.mlp(x)
...         return x.reshape(shape) 

>>> def fix_layer_mlp(layer):
...     layer.old_mlp = layer.mlp
...     layer.mlp = FixedQuantedMLP(layer.mlp)

>>> for layer in base_model.model.layers: fix_layer_mlp(layer)