flash-attention/training/configs/model/gpt2model/gpt2-medium.yaml

# @package _global_
model:
  config:
    n_embd: 1024
    n_head: 16
    n_layer: 24