File size: 3,283 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 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 | """DaisyChain training entry point (CLI: `daisychain-train`).
Reads cluster settings from env (set on each machine, changing only RANK):
MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK -- standard torch.distributed
GLOO_SOCKET_IFNAME -- the NIC to use (e.g. tailscale0)
DAISY_TASK = "module:Class" -- your task (default: example)
DAISY_STEPS = 300
DAISY_LR = 0.05
DAISY_OPTIMIZER = sgd | adam
DAISY_BASE_BATCH = 32
DAISY_STATUS_FILE = status.json -- rank 0 writes live status here
DAISY_STEP_SLEEP = 0 -- demo pacing
DAISY_SAVE = daisychain_model.pt -- rank 0 saves here
"""
import os
import torch
from .cluster import DaisyCluster
from .task import load_task
def _report_verified_counts(cluster):
"""All-reduce verified-unit invocation counts across nodes (if any fired)."""
try:
from .verified import instrument
import torch.distributed as dist
counts = instrument.report()
if not counts:
return
keys = sorted(counts)
t = torch.tensor([counts[k] for k in keys], dtype=torch.float64)
dist.all_reduce(t, op=dist.ReduceOp.SUM)
if cluster.is_master():
print("[verified] CLUSTER-WIDE verified-unit invocations (trained through them):")
for k, v in zip(keys, t.tolist()):
print(f"[verified] {k:34s} {int(v):,}")
except Exception:
pass
def main():
# Default: train THROUGH the emulated GPU logic (verified INT8 units).
# Set DAISY_TASK=daisychain.example_task:ExampleTask for a plain-float run.
task_spec = os.environ.get("DAISY_TASK", "daisychain.verified_task:VerifiedTask")
task = load_task(task_spec)
cluster = DaisyCluster(
cpu_fraction=float(os.environ.get("DAISY_CPU_FRACTION", "0.9")),
base_batch=int(os.environ.get("DAISY_BASE_BATCH", "32")),
)
if cluster.is_master():
p = cluster.plan
print(f"[daisychain] task={task_spec}")
print(f"[plan] world={p['world']} devices={p['devices']} "
f"total_cores={p['total_cores']} total_ram_gb={p['total_ram_gb']}")
print(f"[plan] capacities={p['capacities']} weights={[round(w,3) for w in p['weights']]}")
print(f"[plan] local_batches={p['local_batches']} global_batch={p['global_batch']}")
model = task.build_model()
cluster.fit(
model, task,
steps=int(os.environ.get("DAISY_STEPS", "300")),
lr=float(os.environ.get("DAISY_LR", "0.05")),
optimizer=os.environ.get("DAISY_OPTIMIZER", "sgd"),
status_path=os.environ.get("DAISY_STATUS_FILE", "status.json"),
step_delay=float(os.environ.get("DAISY_STEP_SLEEP", "0")),
)
# if the task trained THROUGH the verified units, report cluster-wide counts
_report_verified_counts(cluster)
diff = cluster.replica_diff(model)
if cluster.is_master():
print(f"[check] replica max param diff across nodes: {diff:.2e}")
save = os.environ.get("DAISY_SAVE", "daisychain_model.pt")
torch.save({"state_dict": model.state_dict()}, save)
print(f"[save] {save}")
cluster.shutdown()
if __name__ == "__main__":
main()
|