This is about the nuts and bolts of training ML models.

• Training/optimization advice/wisdom

  • https://karpathy.medium.com/yes-you-should-understand-backprop-e2f06eab496b
  • https://karpathy.github.io/2019/04/25/recipe/

• Coding mistakes

  • Dimensions (e.g. ytarget not being columnar)
  • Forgetting to zero grads
  • Loss: flipping input / target in MSE, omitting negatives in log loss
  • When evaluating, forgetting model.eval() or no_grad; and vice-versa (forgetting model.train())
  • Forgetting .to(device)

• Bias-variance tradeoff

  

  

• Data processing common things

  • Imputation, missing dummies
  • Imbalanced classes
  • Log skewed data
  • Consolidating sparse columns
  • Standardization

• What to try next

  	Fixes high bias	Fixes high var
  Data		More
  Features	More	Fewer
  Params	More	Fewer
  Regularization	Less	More

• overfitting

  • regularization
  • dropout
  • more data
  • data augmentation
    • train augmentation
    • test time augmentation with voting/mean
  • early stop

• ML projects structure

  • transfer learning vs multi-task learning: transfer if base dataset size is much larger, multi-task best if many tasks where all tasks about same dataset size
  • multi-stage vs end-to-end: just if you have enough data for the subproblems vs the end-to-end problem

• Training issues and optimizations

  • Weight decay: only over weights
  • Vanishing gradients
  • Exploding gradients
  • Mixed precision training

• Dealing with vanishing/exploding gradients

  • initialization
  • norms
  • residual connections

• Debugging activations

  • Visualize distributions of activations:

    • For a simple NN like this:

      layers = [
        Linear(n_embd * block_size, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
        Linear(           n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
        Linear(           n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
        Linear(           n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
        Linear(           n_hidden, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(),
        Linear(           n_hidden, vocab_size, bias=False), BatchNorm1d(vocab_size),
      ]
      

    • Expect these for activations, grads, and weights—don’t want saturation to go to 0, want all layers to be about the same instead of vanishing/exploding:

      plt.figure(figsize=(20, 4)) # width and height of the plot
      legends = []
      for i, layer in enumerate(layers[:-1]): # note: exclude the output layer
        if isinstance(layer, Tanh):
          t = layer.out
          print('layer %d (%10s): mean %+.2f, std %.2f, saturated: %.2f%%' % (i, layer.__class__.__name__, t.mean(), t.std(), (t.abs() > 0.97).float().mean()*100))
          hy, hx = torch.histogram(t, density=True)
          plt.plot(hx[:-1].detach(), hy.detach())
          legends.append(f'layer {i} ({layer.__class__.__name__}')
      plt.legend(legends);
      plt.title('activation distribution')
      

      

      # visualize histograms
      plt.figure(figsize=(20, 4)) # width and height of the plot
      legends = []
      for i, layer in enumerate(layers[:-1]): # note: exclude the output layer
        if isinstance(layer, Tanh):
          t = layer.out.grad
          print('layer %d (%10s): mean %+f, std %e' % (i, layer.__class__.__name__, t.mean(), t.std()))
          hy, hx = torch.histogram(t, density=True)
          plt.plot(hx[:-1].detach(), hy.detach())
          legends.append(f'layer {i} ({layer.__class__.__name__}')
      plt.legend(legends);
      plt.title('gradient distribution')
      

      

      # visualize histograms
      plt.figure(figsize=(20, 4)) # width and height of the plot
      legends = []
      for i,p in enumerate(parameters):
        t = p.grad
        if p.ndim == 2:
          print('weight %10s | mean %+f | std %e | grad:data ratio %e' % (tuple(p.shape), t.mean(), t.std(), t.std() / p.std()))
          hy, hx = torch.histogram(t, density=True)
          plt.plot(hx[:-1].detach(), hy.detach())
          legends.append(f'{i} {tuple(p.shape)}')
      plt.legend(legends)
      plt.title('weights gradient distribution');
      

      

    • Visualize log( SD of grad steps / SD of weights ) for linear layers, over steps:

      ud.append([((lr*p.grad).std() / p.data.std()).log10().item() for p in parameters])
      ...
      plt.figure(figsize=(20, 4))
      legends = []
      for i,p in enumerate(parameters):
        if p.ndim == 2:
          plt.plot([ud[j][i] for j in range(len(ud))])
          legends.append('param %d' % i)
      plt.plot([0, len(ud)], [-3, -3], 'k') # these ratios should be ~1e-3, indicate on plot
      plt.legend(legends);
      

      

• Initialization

  • Wisdom

    • Usu. Gaussian

    • If X, W are unit normal, then XW will have SD > 1, but want SD=1. Should scale init W by $\sqrt{1/d}$ where d is dimension

    • But there’s also the activation nonlinearity, and you want the activation to have SD 1 too.

    • Parameterized Kaiming init is probably the most common init now (and in Pytorch). What gain you use depends on the nonlinearity. E.g. for ReLU, He says init to $\sqrt{2/d}$ (because ReLU discards half the dist).

      

    • GPT-NeoX uses “small init” $\sqrt{\frac{2}{5d}}$ from Transformers without Tears

    • Biases may be OK to zero

    • Residuals should be zeroed

  • Has become less important to init due to modern features like residuals, batch norm/layer norm, and better optimizers (RMSprop, Adam)
  • Resources
    • Learning resources: Karpathy, Howard
    • Papers: Delving Deep Into Rectifiers by He

• Distributed training

  • https://github.com/microsoft/DeepSpeed
  • Torch lightning

• Pretraining

  • one common way to do this atm is to create your own exotic fork of MegatronLM+DeepSpeed. https://news.ycombinator.com/item?id=34421515

• Approaches

  • Siamese network: either 2x for contrastive loss or 3x triplet loss
    • Triplet loss specifically insists there is at least some margin separation between the distances: norm(anchor - pos) - norm(anchor - neg) < -a

• LR schedules

  • Visualizations (source)

    

  • Even though Adam effectively gives each parameter dynamic learning rate, it can still make sense to use a scheduler on what is effectively the global cap of learning rates (source)

  • From https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/guide2/Research_Projects.html

    • A usual good starting point is 0.1 for SGD, and 1e-3 for Adam.
    • The deeper the model is, the lower the learning rate usually should be. For instance, Transformer models usually apply learning rates of 1e-5 to 1e-4 for Adam.
    • For image classifiers and SGD as optimizer, the multi-step LR scheduler has shown to be good choice.
    • Models trained with Adam commonly use a smooth exponential decay in the learning rate or a cosine-like scheduler.
    • For Transformers: remember to use a learning rate warmup. The cosine scheduler is often used for decaying the learning rate afterwards, but can also be replaced by an exponential decay.

• Regularization techniques

  • Weight decay, a form of L2 regularization (if using Adam use AdamW instead)
  • Dropout
  • Augmentation

• Instruction fine tuning

  • follow the 3-step recipe of https://openai.com/blog/chatgpt/ to finetune the model to be an actual assistant instead of just "document completor", which otherwise happily e.g. responds to questions with more questions. Also e.g. see OPT-IML https://arxiv.org/abs/2212.12017 , or BLOOMZ https://arxiv.org/abs/2211.01786 to get a sense of the work involved here.
  • https://news.ycombinator.com/item?id=34421515

• LLM scaling

• Techniques for scaling to larger models than memory

  • https://huggingface.co/docs/transformers/v4.18.0/en/performance

  • Activation checkpointing: rather than materialize all activations for the backprop, recompute some

  • Gradient accumulation

    • Just do multiple smaller batches forward + backward to accumulate gradient before doing an optimization step
    • For example, if you have 4 GPUs and use per_device_train_batch_size=12 and gradient_accumulation_steps=3 you will have an effective batch size of 4123=144.

  • Gradient checkpointing: train models 10x larger than your avail. mem

    • Even when we set the batch size to 1 and use gradient accumulation we can still run out of memory when working with large models. In order to compute the gradients during the backward pass all activations from the forward pass are normally saved. This can create a big memory overhead. Alternatively, one could forget all activations during the forward pass and recompute them on demand during the backward pass. This would however add a significant computational overhead and slow down training. Gradient checkpointing strikes a compromise between the two approaches and saves strategically selected activations throughout the computational graph so only a fraction of the activations need to be re-computed for the gradients. See this great article explaining the ideas behind gradient checkpointing.
    • The slowdown will depend on the model but quite often it is around 20-30%.

  • Mixed precision training: compute (forward+backward) in fp16/bf16 rather than fp32

    • Faster compute (2-8x, typical ~3x) that can leverage certain GPU tensor cores.
    • Must still keep fp32 weights for weight update step. “Although the gradients are also computed in half precision they are converted back to full precision for the optimization step so no memory is saved here. “
    • fp16+fp32 means 1.5x more memory used on weights. Paper notes weights are minority, most memory on activations—but this ignores gradient checkpointing (discarding activations).
    • Loss scaling: many fp16 numbers would round down to 0, so just scale everything up first.
    • BF16: care more about exponent (the range) than the precision

    

    • A cited notebook

  • Low-memory Optimizers: Ada, 8bit Adam

  • Floating point types

    • fp16/bf16 (smaller data/faster throughput)