File size: 1,815 Bytes
309b968 | 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 | # Training your own model
DaisyChain trains any **Task** — an object with three methods:
```python
import torch, torch.nn as nn
class MyTask:
def build_model(self) -> nn.Module:
torch.manual_seed(0) # deterministic -> identical on every node
return nn.Sequential(nn.Linear(16, 64), nn.ReLU(), nn.Linear(64, 10))
def sample(self, n): # this node's data shard
X = torch.randn(n, 16)
y = torch.randint(0, 10, (n,))
return X, y
def loss(self, model, X, y): # mean loss over the batch
return nn.functional.cross_entropy(model(X), y)
```
## Point DaisyChain at it
```bash
export DAISY_TASK="my_task:MyTask" # module:Class, must be importable
daisychain-train
```
Copy `examples/my_task_template.py` to start.
## Rules that matter
1. **`build_model` must be deterministic** (seed it). Every node builds the model
independently, then rank 0's weights are broadcast — but seeding keeps shapes
and buffers consistent.
2. **`sample(n)` should return this node's shard.** For real datasets, split by
`RANK` (e.g. different files/row-ranges per rank) so nodes don't all train on
the same rows. Read `os.environ["RANK"]` / `WORLD_SIZE`.
3. **The model must fit on one node.** DaisyChain pools compute, not memory.
4. Keep it **small.** See [LIMITS.md](LIMITS.md).
## Knobs (env)
| var | default | meaning |
|-----|---------|---------|
| `DAISY_TASK` | example | `module:Class` |
| `DAISY_STEPS` | 300 | training steps |
| `DAISY_LR` | 0.05 | learning rate |
| `DAISY_OPTIMIZER` | sgd | `sgd` or `adam` |
| `DAISY_BASE_BATCH` | 32 | per-node base batch (scaled by capacity) |
| `DAISY_SAVE` | daisychain_model.pt | where rank 0 saves |
| `DAISY_FORCE_CPU` | – | set `1` to ignore a local GPU |
|