Implement decoder-only GPT-style transformer

Goal

Implement a simplified decoder-only Transformer language model (similar in spirit to GPT) for next-token prediction. The implementation should be modular, using four main classes:

1. MultiHeadAttention
2. FeedForward
3. DecoderLayer
4. GPT (the full model)

You may assume a deep learning framework (e.g., PyTorch or TensorFlow), but you should clearly specify tensor shapes and operations.

Model details and requirements

Assume:

• Batch size: B
• Sequence length: T
• Embedding dimension: d_model
• Number of attention heads: num_heads (assume d_model is divisible by num_heads)
• Vocabulary size: V

1. Multi-head self-attention (MultiHeadAttention)

Implement a class MultiHeadAttention that performs masked self-attention:

• Inputs:
• x: Tensor of shape (B, T, d_model) (token embeddings or layer outputs).
• Internally learnable parameters:
• Linear projections to compute queries Q, keys K, and values V for all heads.
• Operations:

1. Project x into Q, K, and V with shapes (B, T, d_model).
2. Split them into num_heads heads, each of dimension d_head = d_model / num_heads.
3. Compute scaled dot-product attention for each head: [ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_{head}}} + M\right)V, ] where M is a causal mask that prevents a position from attending to future positions (positions > t).
4. Concatenate the outputs of all heads back to shape (B, T, d_model).
5. Apply a final linear projection to return to (B, T, d_model).

• Output:
• Tensor of shape (B, T, d_model).

Ensure you correctly implement the causal (autoregressive) mask so that position t can only attend to positions 0..t.

2. Position-wise feed-forward network (FeedForward)

Implement a class FeedForward for the position-wise MLP:

• Inputs:
• x: Tensor of shape (B, T, d_model).
• Internal structure (typical choice):
• Linear layer from d_model to d_ff (e.g., 4 * d_model).
• Non-linear activation (e.g., GELU or ReLU).
• Linear layer from d_ff back to d_model.
• All operations are applied independently to each position in the sequence.
• Output:
• Tensor of shape (B, T, d_model).

3. Decoder layer (DecoderLayer)

Implement a single Transformer decoder block that combines attention, feed-forward, residual connections, and layer normalization:

• Inputs:
• x: Tensor of shape (B, T, d_model).
• Components:

1. LayerNorm ln1 before self-attention.
2. MultiHeadAttention block (masked self-attention).
3. Residual connection: x = x + attention_output.
4. LayerNorm ln2 before feed-forward.
5. FeedForward block.
6. Residual connection: x = x + ff_output.

• Output:
• Tensor of shape (B, T, d_model).

Your DecoderLayer should be reusable so that multiple layers can be stacked.

4. GPT model (GPT)

Implement a GPT class representing the full decoder-only Transformer:

• Inputs:
• input_ids: Integer tensor of shape (B, T) representing token indices.
• Components:

1. Token embedding layer mapping token IDs to vectors of size d_model.
2. Positional encodings (learned or fixed) of shape (T, d_model) added to the token embeddings.
3. A stack of N identical DecoderLayer blocks.
4. A final linear layer mapping from d_model to vocabulary size V.

• Forward pass:

1. Embed tokens, add positional encodings to obtain (B, T, d_model).
2. Pass through the N decoder layers sequentially.
3. Apply the final linear layer to obtain logits of shape (B, T, V).

• Output:
• Logits for next-token prediction: (B, T, V).

Additional notes

• You should define the forward method for each class with correct tensor transformations.
• Be careful about tensor shapes when splitting/combining heads.
• Ensure the causal mask is correctly broadcast and applied in attention.
• You do not need to implement training loops or optimization; focus on correct and clean model implementation.