Add support to Llama 3.1

Nobody has claimed this yet.

### Description & Motivation

In the Tensor Parallel example [here](https://github.com/Lightning-AI/pytorch-lightning/tree/master/examples/fabric/tensor_parallel) while the proposed model is a Llama3 like model, the forward pass is different from the Llama3.1 model in the official code [here](https://github.com/meta-llama/llama3/blob/main/llama/model.py), by example the `Transformers`'s forward pass:

```python
def forward(self, tokens: torch.Tensor, start_pos: int):
_bsz, seqlen = tokens.shape
h = self.tok_embeddings(tokens)
self.freqs_cis = self.freqs_cis.to(h.device)
freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

mask = None
if seqlen > 1:
mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)

mask = torch.triu(mask, diagonal=1)

# When performing key-value caching, we compute the attention scores
# only for the new sequence. Thus, the matrix of scores is of size
# (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for
# j > cache_len + i, since row i corresponds to token cache_len + i.
mask = torch.hstack(
[torch.zeros((seqlen, start_pos), device=tokens.device), mask]
).type_as(h)

for layer in self.layers:
h = layer(h, start_pos, freqs_cis, mask)
h = self.norm(h)
output = self.output(h).float()
return output
```

so there is the additional support to the mask, etc.
This cause issues when trying to train and run the updated Llama3.1 model from that example.
While I have tried to adapt the llama31 code to the example code, removing the mask

```python
def forward(self, tokens: torch.Tensor, start_pos: int):
_bsz, seqlen = tokens.shape
h = self.tok_embeddings(tokens)
self.freqs_cis = self.freqs_cis.to(h.device)
freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]

mask = None
if seqlen > 1:
mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)

mask = torch.triu(mask, diagonal=1)

# When performing key-value caching, we compute the attention scores
# only for the new sequence. Thus, the matrix of scores is of size
# (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for
# j > cache_len + i, since row i corresponds to token cache_len + i.
mask = torch.hstack(
[torch.zeros((seqlen, start_pos), device=tokens.device), mask]
).type_as(h)

for layer in self.layers.values():
#h = layer(h, start_pos, freqs_cis, mask)
h = layer(h, freqs_cis)
h = self.norm(h)
output = self.output(h).float()
return output
```

I'm not sure this is correct in inference infact I 'm getting a number of bad generated tokens from my generation function

```python
from typing import Optional
def text_completion(tokenizer, model, prompts: list, temperature: float = 0.6, top_p: float = 0.9, max_gen_len: Optional = None):
if max_gen_len is None:
max_gen_len = max_seq_len - 1
# Convert each prompt into tokens
prompt_tokens = [tokenizer.encode(prompt, add_bos=True, add_eos=False) for prompt in prompts]
# Make sure the batch size is not too large
batch_size = len(prompt_tokens)
assert batch_size <= max_batch_size, f"batch size must be less than or equal to {max_batch_size}"
max_prompt_len = max(len(prompt) for prompt in prompt_tokens)
# Make sure the prompt length is not larger than the maximum sequence length
assert max_prompt_len <= max_seq_len, f"prompt length must be less than or equal to {max_seq_len}"
total_len = min(max_seq_len, max_gen_len + max_prompt_len)

# Create the list that will contain the generated tokens, along with the initial prompt tokens
pad_id = tokenizer.pad_id
tokens = torch.full((batch_size, total_len), pad_id, dtype=torch.long, device=device)
for k, t in enumerate(prompt_tokens):
# Populate the initial tokens with the prompt tokens
tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device=device)

eos_reached = torch.tensor([False] * batch_size, device=device)
prompt_tokens_mask = tokens != pad_id # True if the token is a prompt token, False otherwise
cur_iterator = tqdm(range(1, total_len), desc="Generating tokens")
for cur_pos in cur_iterator:
with torch.no_grad():
logits = model.forward(tokens[:, cur_pos-1:cur_pos], cur_pos)
#logits = model.forwar(tokens[:, cur_pos-1:cur_pos])
if temperature > 0:
# The temperature is applied before the softmax
probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
next_token = _sample_top_p(probs, top_p)
else:
# Greedily select the token with the max probability
next_token = torch.argmax(logits[:, -1], dim=-1)

next_token = next_token.reshape(-1)
# Only replace token if it is a padding token
next_token = torch.where(prompt_tokens_mask[:, cur_pos], tokens[:, cur_pos], next_token)
tokens[:, cur_pos] = next_token
# EOS is reached only if we found an EOS token for a padding position
eos_reached |= (~prompt_tokens_mask[:, cur_pos]) & (next_token == tokenizer.eos_id)
if all(eos_reached):
break

out_tokens = []
out_text = []
for prompt_index, current_prompt_tokens in enumerate(tokens.tolist()):
# Cut to the EOS token, if present
if tokenizer.eos_id in current_prompt_tokens:
eos_idx = current_prompt_tokens.index(tokenizer.eos_id)
current_prompt_tokens = current_prompt_tokens[:eos_idx]
out_tokens.append(current_prompt_tokens)
out_text.append(tokenizer.decode(current_prompt_tokens))
return (out_tokens, out_text)
```

the output token emitted are `eos`, the input token from the string `The` and a sequence of `0` padded token

```
Generating tokens: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 127/127 [00:04<00:00, 31.37it/s]
([[128000, 791, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], ['The!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'])

```

### Pitch

Add full support to Llama3 updated model code from meta-llama [here](https://github.com/meta-llama/llama3/blob/main/llama/model.py)
This issue is also referenced here https://github.com/meta-llama/llama/issues/306

### Alternatives

adapt the current code to support masking

### Additional context

_No response_

cc @borda

Begin with examples/fabric/tensor_parallel and compare its model implementation with Meta's llama/model.py and the forward and generation snippets in the issue. Check masking and cache behavior against the referenced Llama 3.1 implementation, then run the example's training and inference path; done means Llama 3.1 runs without the reported EOS, repeated-token, and padding output.

Written by the indexing model from the issue text.