fix DDP hang: spawn dataloader workers (fork-after-CUDA/NCCL starved torchrun heartbeat), picklable Collate, fewer workers per rank

tinyvla_b200/scripts/train_fast.py

@@ -96,6 +96,58 @@ def load_compatible(model, path: Path):
     return len(keep)
 
 
+
+class Collate:
+    """Picklable collate: tokenizes task strings with a per-process cache.
+
+    A module-level class (not a closure) so DataLoader workers can be started
+    with the 'spawn' context. spawn matters under DDP: forking workers from a
+    process that already initialized CUDA/NCCL (with ffmpeg + rust-tokenizer
+    threads alive) is exactly the fork-after-CUDA hazard that hung rank
+    dataloaders in testing; spawned workers start clean.
+    """
+
+    def __init__(self, model_name: str, max_length: int):
+        self.model_name = model_name
+        self.max_length = max_length
+        self._tokenizer = None
+        self._cache: dict = {}
+
+    def _tok(self, tasks):
+        if self._tokenizer is None:
+            import os
+
+            os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        new = [t for t in set(tasks) if t not in self._cache]
+        if new:
+            e = self._tokenizer(new, padding="max_length", truncation=True,
+                                max_length=self.max_length, return_tensors="pt")
+            for i, t in enumerate(new):
+                self._cache[t] = (e["input_ids"][i], e["attention_mask"][i].bool())
+        return (torch.stack([self._cache[t][0] for t in tasks]),
+                torch.stack([self._cache[t][1] for t in tasks]))
+
+    def __getstate__(self):
+        return {"model_name": self.model_name, "max_length": self.max_length}
+
+    def __setstate__(self, st):
+        self.__init__(st["model_name"], st["max_length"])
+
+    def __call__(self, items):
+        out = {}
+        for k in items[0]:
+            if k == "task":
+                ids, mask = self._tok([it["task"] for it in items])
+                out["observation.language.tokens"] = ids
+                out["observation.language.attention_mask"] = mask
+            else:
+                out[k] = torch.stack([it[k] for it in items])
+        return out
+
+
 # ----------------------------------------------------------------------- main
 
 
@@ -205,53 +257,23 @@ def main():
         )
 
     # ---- pre-tokenize every distinct instruction once -----------------------
-    from transformers import AutoTokenizer
-
-    tokenizer = AutoTokenizer.from_pretrained(pcfg.lm_model_name)
-    if src_cfg.get("type") == "hub":
-        _tok_cache: dict = {}
-
-        def _tok(tasks):
-            new = [t for t in set(tasks) if t not in _tok_cache]
-            if new:
-                e = tokenizer(new, padding="max_length", truncation=True,
-                              max_length=pcfg.tokenizer_max_length, return_tensors="pt")
-                for i, t in enumerate(new):
-                    _tok_cache[t] = (e["input_ids"][i], e["attention_mask"][i].bool())
-            return (torch.stack([_tok_cache[t][0] for t in tasks]),
-                    torch.stack([_tok_cache[t][1] for t in tasks]))
-    else:
-        texts = sorted(set(source._tasks.values()))
-        enc = tokenizer(texts, padding="max_length", truncation=True,
-                        max_length=pcfg.tokenizer_max_length, return_tensors="pt")
-        tok_ids, tok_mask = enc["input_ids"], enc["attention_mask"].bool()
-        tok_lookup = {t: i for i, t in enumerate(texts)}
-        print(f"pre-tokenized {len(texts)} instructions at fixed length {pcfg.tokenizer_max_length}")
-
-        def _tok(tasks):
-            idx = torch.tensor([tok_lookup.get(t, 0) for t in tasks])
-            return tok_ids[idx], tok_mask[idx]
-
-    def collate(items):
-        out = {}
-        for k in items[0]:
-            if k == "task":
-                ids, mask = _tok([it["task"] for it in items])
-                out["observation.language.tokens"] = ids
-                out["observation.language.attention_mask"] = mask
-            else:
-                out[k] = torch.stack([it[k] for it in items])
-        return out
+    collate = Collate(pcfg.lm_model_name, pcfg.tokenizer_max_length)
 
+    nw = cfg.get("num_workers", 12)
+    if ddp:
+        # 16 воркеров/ранг x N рангов душат CPU и heartbeat torchrun-агента
+        nw = cfg.get("num_workers_per_rank", max(4, nw // world))
+        log(f"DDP: {nw} dataloader workers per rank")
     loader = torch.utils.data.DataLoader(
         source,
         batch_size=cfg["batch_size"],
-        num_workers=cfg.get("num_workers", 12),
+        num_workers=nw,
         pin_memory=True,
-        persistent_workers=cfg.get("num_workers", 12) > 0,
-        prefetch_factor=cfg.get("prefetch_factor", 6) if cfg.get("num_workers", 12) > 0 else None,
+        persistent_workers=nw > 0,
+        prefetch_factor=cfg.get("prefetch_factor", 6) if nw > 0 else None,
         drop_last=True,
         collate_fn=collate,
+        multiprocessing_context="spawn" if nw > 0 else None,
     )
 
     # ---- optimizer: backbone at a lower lr, exactly as train.py does --------