Extending Components¶
Coming Soon
Heyoooo so I ran out of time writing docs as I have to do actual machine learning. I'll eventually catch this up but as of right now here's my friend gpt-6-astra who will do the talking.
Choose the class closest to the behavior you want to change. These guides cover base classes; the API reference lists their concrete implementations.
Trainer Extension¶
Extend a trainer to choose its model, data, optimizer, schedule, evaluations, and analyses, or change its forward computation.
Extension Points¶
| Change | Inherit from | Guide |
|---|---|---|
| Define a configurable model building block | Module |
Module contract |
| Change embeddings, layer execution, logits, or loss | GPT |
GPT |
| Change sublayers and residual connections | Block |
Transformer blocks |
| Change Q/K/V projections, masks, or attention computation | SelfAttention |
Attention |
| Change the feed-forward computation | MLP |
MLP |
| Change expert selection or expert implementations | MoE |
Mixture of experts |
Reusable Modules¶
Use these inside your implementation. Links go directly to the API reference.
| Building block | Reference |
|---|---|
| Rotary position encoding | RotaryPosEncoding |
| Multimodal rotary position encoding | MRotaryPosEncoding |
| Layer normalization | LayerNorm |
RMS normalization, gated RMS normalization, and rms() |
RMSNorm and helpers |
| SwiGLU activation | swiglu |