Merge #469: COLI_MODEL_DIRS — split model shards across N drives (capacity aggregation)
Search-path of extra directories, each holding a distinct subset of the .safetensors shards (no duplication). Each shard lives on one drive; demand preads hit whichever drive holds it, so concurrent expert loads parallelise across drives and combined capacity is used — a 400 GB container fits across two 250 GB drives the mirror (#421) can't. st_init kept as a back-compat wrapper. Verified locally: clean build, no regression on the default path (token-exact tiny models unchanged), and the split path itself validated token-exact (shard on a separate dir, metadata in primary -> 32/32 vs oracle). Composable with the #421 mirror. CI 8/8. Thanks @mohamedmastouri2000-boop.
This commit is contained in:
Vincenzo Fornaro
+3
-1
@@ -1031,7 +1031,9 @@ static void layer_cuda_shard_kvb(Layer *l,int H,int Q,int V){
|
||||
|
||||
static void model_init(Model *m, const char *snap, int cap, int ebits, int dbits){
|
||||
memset(m,0,sizeof(*m)); m->ebits=ebits; m->dbits=dbits;
|
||||
load_cfg(&m->c,snap); st_init(&m->S,snap);
|
||||
load_cfg(&m->c,snap);
|
||||
{ const char *xd=getenv("COLI_MODEL_DIRS"); /* SPLIT: model shards spread across N drives */
|
||||
st_init_multi(&m->S,snap,(xd&&*xd)?xd:NULL); }
|
||||
Cfg *c=&m->c; char nm[256]; int H=c->n_heads, D=c->hidden;
|
||||
/* embed e lm_head sono il confine I/O: tenerli ad alta precisione (come i quant dynamic
|
||||
* reali). A bf16 ~1.9GB su GLM reale: trascurabile. dbits>=8 -> qui f32; piu' basso -> dbits. */
|
||||
|
||||
@@ -210,21 +210,66 @@ static void st_pread_full(int fd, void *buf, int64_t n, int64_t off, const char
|
||||
}
|
||||
}
|
||||
|
||||
static void st_init(shards *S, const char *snap_dir) {
|
||||
/* Scan one directory for *.safetensors shards, appending to files[] (dedup by
|
||||
* basename, so a list of directories acts as a SEARCH PATH: the same shard
|
||||
* present on two drives is taken from the first-listed one only). *added
|
||||
* returns how many shards this dir contributed. */
|
||||
static void st_scan_dir(const char *dir, char files[][1024], int *nf, int *added) {
|
||||
DIR *d = opendir(dir); struct dirent *e;
|
||||
if (!d) { perror(dir); exit(1); }
|
||||
int base_n = *nf;
|
||||
while ((e = readdir(d))) {
|
||||
const char *dot = strrchr(e->d_name, '.');
|
||||
if (dot && !strcmp(dot, ".safetensors")) { /* model.safetensors o model-0000N-of-... */
|
||||
int dup = 0;
|
||||
for (int i = 0; i < *nf; i++) {
|
||||
const char *b = strrchr(files[i], '/');
|
||||
#ifdef _WIN32
|
||||
const char *b2 = strrchr(files[i], '\\'); if (b2 && (!b || b2 > b)) b = b2;
|
||||
#endif
|
||||
b = b ? b + 1 : files[i];
|
||||
if (!strcmp(b, e->d_name)) { dup = 1; break; } /* already taken from a higher-priority drive */
|
||||
}
|
||||
if (dup) continue;
|
||||
if (*nf >= ST_MAX_SHARDS) { fprintf(stderr, "too many shards (>%d): raise ST_MAX_SHARDS\n", ST_MAX_SHARDS); exit(1); }
|
||||
snprintf(files[(*nf)++], 1024, "%s/%s", dir, e->d_name);
|
||||
}
|
||||
}
|
||||
closedir(d);
|
||||
if (added) *added = *nf - base_n;
|
||||
}
|
||||
|
||||
/* Index shards from snap_dir, optionally SPLIT across extra drives listed in
|
||||
* extra_dirs (';' or ',' separated). Each shard lives on exactly ONE drive
|
||||
* (no duplication — unlike the dual-SSD mirror); a demand pread hits whichever
|
||||
* drive holds that shard, so concurrent expert loads parallelise across drives
|
||||
* and combined capacity is used. Scales to N drives. Metadata (config /
|
||||
* tokenizer / .coli_usage / .coli_kv) is read from snap_dir only. */
|
||||
static void st_init_multi(shards *S, const char *snap_dir, const char *extra_dirs) {
|
||||
memset(S, 0, sizeof(*S));
|
||||
S->cap = 4096; S->t = calloc(S->cap, sizeof(st_tensor));
|
||||
/* raccoglie ordinatamente i nomi dei file shard */
|
||||
static char files[ST_MAX_SHARDS][1024]; int nf = 0;
|
||||
DIR *d = opendir(snap_dir); struct dirent *e;
|
||||
if (!d) { perror(snap_dir); exit(1); }
|
||||
while ((e = readdir(d))) {
|
||||
const char *dot = strrchr(e->d_name, '.');
|
||||
if (dot && !strcmp(dot, ".safetensors")) { /* model.safetensors o model-0000N-of-... */
|
||||
if (nf >= ST_MAX_SHARDS) { fprintf(stderr, "too many shards (>%d): raise ST_MAX_SHARDS\n", ST_MAX_SHARDS); exit(1); }
|
||||
snprintf(files[nf++], 1024, "%s/%s", snap_dir, e->d_name);
|
||||
int c0 = 0; st_scan_dir(snap_dir, files, &nf, &c0);
|
||||
int ndir = 1;
|
||||
if (extra_dirs && *extra_dirs) {
|
||||
char buf[4096]; snprintf(buf, sizeof(buf), "%s", extra_dirs);
|
||||
char *p = buf;
|
||||
while (p && *p) {
|
||||
char *sep = p; while (*sep && *sep != ';' && *sep != ',') sep++;
|
||||
int last = (*sep == 0); *sep = 0;
|
||||
while (*p == ' ') p++;
|
||||
size_t plen = strlen(p); while (plen > 0 && p[plen-1] == ' ') p[--plen] = 0;
|
||||
if (*p) {
|
||||
int cN = 0; st_scan_dir(p, files, &nf, &cN);
|
||||
fprintf(stderr, "[SPLIT] +%s -> %d shard(s)\n", p, cN);
|
||||
ndir++;
|
||||
}
|
||||
p = last ? NULL : sep + 1;
|
||||
}
|
||||
fprintf(stderr, "[SPLIT] model across %d dir(s): %d shard(s) total (primary %s -> %d shard(s)), no duplication\n",
|
||||
ndir, nf, snap_dir, c0);
|
||||
}
|
||||
closedir(d);
|
||||
for (int a = 0; a < nf; a++) for (int b = a+1; b < nf; b++)
|
||||
if (strcmp(files[a], files[b]) > 0) { char tmp[1024]; strcpy(tmp, files[a]); strcpy(files[a], files[b]); strcpy(files[b], tmp); }
|
||||
|
||||
@@ -305,6 +350,9 @@ static void st_init(shards *S, const char *snap_dir) {
|
||||
}
|
||||
}
|
||||
|
||||
/* backward-compatible single-directory entry point */
|
||||
static void st_init(shards *S, const char *snap_dir) { st_init_multi(S, snap_dir, NULL); }
|
||||
|
||||
static st_tensor *st_find(shards *S, const char *name) {
|
||||
if (S->hidx) {
|
||||
uint64_t h = st_hash(name) & (S->hcap - 1);
|
||||
|
||||
@@ -82,6 +82,7 @@ Format: `VAR` — default — effect.
|
||||
|
||||
| Variable | Default | Effect |
|
||||
|---|---|---|
|
||||
| `COLI_MODEL_DIRS` | unset | SPLIT the model across 2+ drives: a `;`/`,`-separated list of extra directories, each holding a **distinct** subset of the `.safetensors` shards (no duplication). Shards act as a search path — every shard is read from whichever drive holds it, so concurrent expert loads parallelise across drives and combined capacity is used. Scales to N drives. Metadata (config/tokenizer/`.coli_usage`) stays in the primary `COLI_MODEL` dir. Pairs well with `PIPE=1` (concurrent loaders) + `DIRECT=1`. Distinct from — and composable with — `COLI_MODEL_MIRROR`: the mirror is matched per-shard by basename against the merged (split) index, so a mirror dir may hold a copy of any subset of the split's shards. |
|
||||
| `COLI_MODEL_MIRROR` | unset | Path to a second, byte-identical (read-only) copy of the model on another drive; expert reads are split across both. Partial mirrors work (only the shards present are used). |
|
||||
| `COLI_DISK_WEIGHTS` | unset (startup bandwidth probe) | Split ratio `<primary>,<mirror>` (e.g. `1,1` for 50/50, `9,3` for a fast+slow pair). Unset = probe both drives with the engine's own access pattern at startup. |
|
||||
|
||||
|
||||
Reference in New Issue
Block a user