File size: 6,163 Bytes
a181ec9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | # Adding New Speculative Decoding Algorithms
This guide explains how to add a new speculative decoding algorithm to the Speculators library.
## Quick Start
Adding a new algorithm requires:
1. **Create algorithm module** under `src/speculators/models`.
2. **Configuration class** with `@register` decorator. When Python imports your module, the `@register("myalgo")` decorator adds your class to a global registry dictionary. The training script looks up `"myalgo"` in the registry to find your class. This is helpful because the training script doesn't need to know about every algorithm and adding a new algorithm doesn't require modifying the training script.
3. **Model class** with `@register` decorator
4. **Training factory methods** as classmethods on the model
5. **CLI arguments** in `src/speculators/train/cli.py`
## Step-by-Step Guide
### 1. Create Algorithm Module
Create a self-contained directory for your algorithm under `src/speculators/models`. See `src/speculators/models/eagle3` as an example. This keeps algorithm logic isolated and maintainable. Each algorithm owns its configuration, model definition, and any custom components. Example file structure:
```
src/speculators/models
|-> eagle3
|-> ...
|-> new_algorithm
|-> __init__.py
|-> core.py
|-> config.py
```
### 2. Implement Configuration Class
Define how your algorithm is configured. The config stores hyperparameters, architectural choices, and other settings. It's serialized when saving models and deserialized when loading them. In `config.py`, create a configuration class with the `@register` decorator, for example:
```python
from speculators import SpeculatorModelConfig
@SpeculatorModelConfig.register("myalgo")
class MyAlgoSpeculatorConfig(SpeculatorModelConfig):
speculators_model_type: str = "myalgo"
# Algorithm-specific parameters
block_size: int = 8
num_layers: int = 1
```
**Reference:** See `src/speculators/models/eagle3/config.py` for a complete example.
**Key points:**
- Use `@SpeculatorModelConfig.register("myalgo")` decorator
- Set `speculators_model_type` to match your algorithm name
- Inherit common fields from `SpeculatorModelConfig`
- Add algorithm-specific parameters as needed
### 3. Implement Model Class
Define your algorithm's architecture and training interface. The model class contains model architecture, forward pass logic, and training setup. By implementing the required methods, your algorithm should work seamlessly with the training infrastructure.
In `core.py`, create a model class with the `@register` decorator and required training factory methods.
**Reference:** See `src/speculators/models/eagle3/core.py` for a complete example.
**Required for the training infrastructure:**
Model attributes:
- `layers`: ModuleList of decoder layers (each layer is individually wrapped by FSDP for distributed training)
Methods:
- `from_training_args(cls, verifier_config, **kwargs)`: Factory method to build from CLI args (receives all args as kwargs)
- `get_trainer_kwargs(**kwargs)`: Returns `(train_kwargs, val_kwargs)` dicts passed to `forward()`
- `forward(...)`: Must return `(output, loss, metrics)` where metrics includes a `"loss"` key
### 4. Export Classes
Make your classes importable from the package. Python's import system requires explicit exports from `__init__.py`. This also provides a clean public API.
In `__init__.py`, export your config and model classes.
```python
from speculators.models.eagle3.config import Eagle3SpeculatorConfig
from speculators.models.eagle3.core import Eagle3DraftModel
__all__ = [
"Eagle3DraftModel",
"Eagle3SpeculatorConfig",
]
```
**Reference:** See `src/speculators/models/eagle3/__init__.py`
### 5. Add CLI Arguments (Optional)
Add algorithm-specific command-line arguments to the training script. If your algorithm has unique hyperparameters (like Eagle3's `--ttt-steps` or a custom `--block-size`), users need a way to configure them from the command line. These arguments are passed to your `from_training_args()` method. Only add arguments if your algorithm needs parameters beyond the common ones (verifier path, number of layers, etc.).
**Reference:** See `src/speculators/train/cli.py`
### 6. Train Your Model
The training script should automatically works with your new algorithm:
```bash
torchrun --nnodes=1 --nproc_per_node=8 -m speculators.train \
--speculator-type myalgo \
--verifier-name-or-path meta-llama/Llama-3.1-8B \
--num-layers 1 \
--block-size 8 \
--data-path ./output \
--save-path ./output/checkpoints \
--epochs 20
```
## How It Works
**The flow during training:**
1. User runs: `speculators train --speculator-type myalgo`
2. Training script calls: `model_class = SpeculatorModel.get_class("myalgo")`
3. Registry returns: `MyAlgoDraftModel` class
4. Script converts args to dict: `vars(args)` and calls: `model_class.from_training_args(verifier_config, **vars(args))`
5. Your factory method extracts the kwargs it needs and builds the model instance
6. Trainer validates the model is registered (via checks in `setup_model()` and `apply_fully_sharded()`)
This pattern is similar to how `transformers` uses `.from_pretrained()` - each model owns its own instantiation logic.
**Reference:** See `src/speculators/train/cli.py`
## Using Base Components
Shared transformer layer components that can be reused across algorithms. Many speculative decoding algorithms use similar architectural components (decoder layers, attention, normalization). Instead of duplicating code, you can import pre-configured components for different base model architectures.
**When to use:** If your algorithm uses standard transformer components from models like LLaMA or Qwen3, you can import them from `base_components` instead of defining your own. This is especially useful when you only need to customize one layer (like the first layer) while keeping the rest standard.
**Available architectures:** `llama`, `qwen3`
**Reference:**
- Component definitions: `src/speculators/models/base_components.py`
- Usage example: `src/speculators/models/eagle3/model_definitions.py`
|