|
| 1 | +# Adding a new language model architecture to XLM |
| 2 | + |
| 3 | +For scaffolding a new external model, see [External Models](external-models.md). This document outlines the process for adding a new language model architecture to the XLM codebase. The code follows a modular design with four main components that work together to provide a complete language modeling solution. You need to implement all of them in order to add a new working model. |
| 4 | + |
| 5 | +## Overview of Main Components |
| 6 | + |
| 7 | +### 1. **LossFunction** |
| 8 | +The `LossFunction` is responsible for computing the training loss during model training, validation and optionally test time. |
| 9 | + |
| 10 | +**Key Responsibilities:** |
| 11 | +- Compute loss between model predictions and ground truth targets |
| 12 | +- Return a dictionary with "loss" key and any other additional values that you want to track. |
| 13 | + |
| 14 | +**Interface:** |
| 15 | +```python |
| 16 | +class LossFunction(Generic[T_in, T_out], Protocol): |
| 17 | + model: Any |
| 18 | + tokenizer: Tokenizer |
| 19 | + |
| 20 | + def loss_fn(self, batch: T_in, ...) -> T_out: ... |
| 21 | + def configure(self, pl_module: "Harness"): ... |
| 22 | + """Converts scalar to tensor such that loss_fn becomes compile friendly. If you don't want to compile you don't need to implement this. |
| 23 | + """ |
| 24 | +``` |
| 25 | + |
| 26 | +Examples: `xlm-models/arlm/loss_arlm.py` |
| 27 | + |
| 28 | +### 2. **Predictor** |
| 29 | +The `Predictor` handles generating output sequences from the model. |
| 30 | + |
| 31 | +**Key Responsibilities:** |
| 32 | +- Run the model (typically in a loop) to produce a sequence of tokens. |
| 33 | +- Convert generated token_ids to text. |
| 34 | + |
| 35 | +**Interface:** |
| 36 | +```python |
| 37 | +class Predictor(Generic[T_in, T_out_pred], Protocol): |
| 38 | + tokenizer: Tokenizer |
| 39 | + noise_schedule: NoiseSchedule |
| 40 | + model: Any |
| 41 | + |
| 42 | + def predict(self, batch: T_in, ...) -> T_out_pred: ... |
| 43 | + def to_dict(self, batch: T_in, preds: T_out_pred, ...) -> List[Dict[str, Any]]: ... |
| 44 | +``` |
| 45 | + |
| 46 | +Examples: `xlm-models/arlm/predictor_arlm.py` |
| 47 | + |
| 48 | +### 3. **Collator** |
| 49 | +The `Collator` prepares batches of data for training and inference. It handles data preprocessing, padding, and batching. |
| 50 | + |
| 51 | +**Key Responsibilities:** |
| 52 | +- It receives raw token_ids and converts them to a dict which is passed in as a batch to the model. |
| 53 | +- Handle padding and truncation. |
| 54 | + |
| 55 | +**Interface:** |
| 56 | +```python |
| 57 | +class Collator(Protocol): |
| 58 | + """For pre-training the model on language modeling.""" |
| 59 | + def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, Any]: ... |
| 60 | + |
| 61 | +class Seq2SeqCollator(Collator): |
| 62 | + """For training the model on seq2seq tasks.""" |
| 63 | + def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, Any]: ... |
| 64 | + |
| 65 | +class Seq2SeqCollatorPrediction(Collator): |
| 66 | + """For generating predictions for seq2seq tasks.""" |
| 67 | + def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, Any]: ... |
| 68 | +``` |
| 69 | + |
| 70 | +Examples: `xlm-models/arlm/datamodule_arlm.py` |
| 71 | + |
| 72 | +### 4. **Model** |
| 73 | +The `Model` is the bare neural network architecture for your LM. It defines the forward pass and model parameters. |
| 74 | + |
| 75 | +**Key Responsibilities:** |
| 76 | +- Define the neural network architecture. |
| 77 | +- Implement the forward pass. |
| 78 | + |
| 79 | +**Interface:** |
| 80 | +```python |
| 81 | +class Model: |
| 82 | + def forward(self, input_ids: Tensor, attention_mask: Optional[Tensor] = None, ...) -> Tensor: ... |
| 83 | +``` |
| 84 | + |
| 85 | +Examples: `xlm-models/arlm/model_arlm.py` |
| 86 | + |
| 87 | +All these four components are designed to be aware of each other, and are only expected to run with each other for the same LM and not with any other LM. This is a key design choice that allows one to implement really esoteric models without worrying about how to abstract them such that their dataflow becomes compatible with other LMs. |
| 88 | + |
| 89 | +## Step-by-Step Guide |
| 90 | + |
| 91 | +See the [External Models](external-models.md) guide for the recommended workflow using `xlm-scaffold`. The scaffold creates the directory structure, type definitions, and config stubs. You then implement the four components (Model, Loss, Predictor, Collator) and register them in Hydra configuration. |
| 92 | + |
| 93 | +Key config locations (paths may vary for external models in `xlm-models/`): |
| 94 | + |
| 95 | +- `configs/lightning_train/collator/` – Collator configs |
| 96 | +- `configs/lightning_train/datamodule/` – Datamodule configs |
| 97 | +- `configs/lightning_train/model/` – Model (neural network) configs |
| 98 | +- `configs/lightning_train/model_type/` – Loss, predictor, metrics |
| 99 | +- `configs/lightning_train/experiment/` – Experiment configs |
| 100 | + |
| 101 | +## Design |
| 102 | + |
| 103 | +### Why is there so much nesting in the datamodule config? |
| 104 | + |
| 105 | +The main component of the datamodule are the `dataset_managers`. Each `dataset_manager` will generate its own `dataloader` with its own `collator` and processing functions. This design allows: |
| 106 | + |
| 107 | +1. Chaining arbitrary number of "eval" tasks/datasets during validation or testing |
| 108 | +2. Injecting new eval tasks post-training |
| 109 | + |
| 110 | +## Best Practices |
| 111 | + |
| 112 | +- **Type Safety**: Use `jaxtyping` for tensor type annotations; define clear interfaces with `TypedDict` and `Protocol` |
| 113 | +- **Modularity**: Keep components loosely coupled; use dependency injection through configuration |
| 114 | +- **Testing**: Use debug mode configs (`debug=overfit`, `debug=small_data`) for quick testing |
| 115 | +- **Documentation**: Document all public interfaces; include examples in docstrings |
| 116 | + |
| 117 | +## Troubleshooting |
| 118 | + |
| 119 | +1. **Hydra Errors**: If you see `Unable to find or instantiate abc.xyz.MyClass`, try importing manually: `python -c "from abc.xyz import MyClass"`. |
| 120 | +2. **Unable to implement a component**: Check existing models in `xlm-models/` (arlm, mlm, ilm, mdlm) for reference. |
0 commit comments