NVIDIA: Deep Learning Engineer

import torch, torch.nn as nn, math

class MHSA(nn.Module):
    def __init__(self, d_model, n_heads, causal=False):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.dk, self.causal = n_heads, d_model // n_heads, causal
        self.qkv = nn.Linear(d_model, 3 * d_model); self.out = nn.Linear(d_model, d_model)
    def forward(self, x):                       # x: (B, T, C)
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
        q, k, v = (t.view(B, T, self.h, self.dk).transpose(1, 2) for t in (q, k, v))   # (B,h,T,dk)
        att = q @ k.transpose(-2, -1) / math.sqrt(self.dk)
        if self.causal:
            mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
            att = att.masked_fill(mask, float("-inf"))
        y = att.softmax(-1) @ v                 # (B,h,T,dk)
        return self.out(y.transpose(1, 2).reshape(B, T, C))

• Discuss the O(T²) memory, FlashAttention (tiled, IO-aware, no T×T matrix in HBM), F.scaled_dot_product_attention, and GQA/MQA for a smaller KV cache.

• Data parallelism (DDP): each GPU holds a full model copy and processes a different batch; gradients are all-reduced. FSDP/ZeRO also shards the parameters, gradients and optimiser states.
• Tensor parallelism (Megatron-LM): split individual weight matrices across GPUs.
• For an MLP, split the first linear layer column-wise and the second row-wise, so only one all-reduce is needed per block.
• Attention heads are split across GPUs.
• It needs fast interconnect (NVLink), so it is usually used within a node.
• Pipeline parallelism: split layers into stages across GPUs, with micro-batches to reduce "bubbles" (GPipe, 1F1B).
• Sequence / context parallelism splits long sequences. Large LLM training combines all of these (3D/4D parallelism).

• Without non-linearity, stacked layers collapse into a single linear map, so activations give networks their expressive power.
• Sigmoid/tanh: saturate and cause vanishing gradients; sigmoid is now mostly used for gates and outputs.
• ReLU: cheap and non-saturating for x > 0, but can produce dead units.
• GELU / SiLU (Swish): smooth, used in Transformers.
• SwiGLU (gated) is used in modern LLMs (Llama) for better quality per FLOP.

• Weight decay (L2; decoupled in AdamW), dropout, data augmentation (mixup, cutmix, random erasing), label smoothing, early stopping, stochastic depth.
• Normalisation (batch/layer norm) has a regularising side effect.
• For LLM fine-tuning: a low learning rate, LoRA, fewer epochs, mixing in general data.

• Mixed precision: compute in FP16/BF16 on Tensor Cores (2–4× faster, half the memory) with FP32 master weights.
• FP16 needs loss scaling to avoid underflow.
• BF16 has the FP32 exponent range, so it is more stable.
• FP8 (Hopper/Blackwell, Transformer Engine) adds per-tensor scaling for further speed.
• Gradient (activation) checkpointing: don't store all activations; recompute them in the backward pass. This trades about 30% extra compute for large memory savings, so you can use bigger batches or models.

• Pre-norm and RMSNorm for stability; RoPE positional encoding; SwiGLU feed-forward layers.
• GQA/MQA for an efficient KV cache; FlashAttention kernels; sliding-window or sparse attention for long context.
• Mixture-of-Experts (Mixtral, DeepSeek) for more parameters at a similar compute per token.
• Inference: speculative decoding, paged KV cache (vLLM), quantization (FP8/INT4).
• Training: scaling laws, better data mixtures, RLHF, DPO and RL with verifiable rewards for reasoning.