docs/CUSTOM_TASK.md

Training your own model

DaisyChain trains any Task — an object with three methods:

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

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.

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