fix: serialize model/tokenizer loading across ranks to avoid HF cache contention

configs/trl/dpo.yaml

@@ -20,6 +20,7 @@ method: dpo
 backend: trl
 run_name: null                               # auto-generated from model + datasets if null
 offline: false                               # set true to disable all HuggingFace / wandb network calls
+load_model_serially_across_ranks: false      # setting to true avoids HF cache lock contention on some parallel filesystems (e.g. Lustre)
 
 # -- Model -------------------------------------------------------------------
 model:

configs/trl/sft.yaml

@@ -20,6 +20,7 @@ method: sft
 backend: trl
 run_name: null                               # auto-generated from model + datasets if null
 offline: true
+load_model_serially_across_ranks: false      # setting to true avoids HF cache lock contention on some parallel filesystems (e.g. Lustre)
 
 # Container (set to null, remove, or set image: null for bare-metal)
 container:

src/post_training/config.py

@@ -236,6 +236,7 @@ class PostTrainingConfig:
     llamafactory: dict | None = None
     container: ContainerConfig | None = None
     prefetch_assets: bool = True
+    load_model_serially_across_ranks: bool = False
 
     model: ModelConfig = field(default_factory=ModelConfig)
     training: TrainingConfig = field(default_factory=TrainingConfig)

src/post_training/methods/common.py

@@ -9,10 +9,13 @@
 import dataclasses
 import logging
 import os
+from collections.abc import Callable
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
+import torch.distributed as dist
+from accelerate import PartialState
 from transformers import AutoTokenizer
 
 from post_training.callbacks.inference_checkpoint import InferenceCheckpointCallback
@@ -25,6 +28,66 @@
 
 logger = logging.getLogger(__name__)
 
+
+def build_one_at_a_time[T](
+    state: PartialState, build_fn: Callable[[], T], serial: bool = True
+) -> T:
+    """Call ``build_fn`` on exactly one process at a time, in rank order.
+
+    ``from_pretrained`` resolves the HF Hub cache via ``filelock``-guarded
+    reads/writes that are unreliable under N-way concurrent access on some
+    parallel filesystems (observed on Lustre: spurious "does not appear to
+    have a file named ..." errors on an already fully-cached, static
+    checkpoint, triggered only when every rank calls ``from_pretrained``
+    simultaneously). Serializing via a rank-ordered barrier avoids any two
+    processes contending for the same lock at once, without depending on
+    that same filesystem for the ordering guarantee itself (the barrier
+    goes over the NCCL/Gloo interconnect).
+
+    Set ``serial=False`` (``load_model_serially_across_ranks: false`` in the
+    config) to skip the barrier and let every rank call ``build_fn``
+    concurrently, e.g. on filesystems that don't exhibit this contention.
+
+    If ``build_fn`` raises on the active rank, that rank is caught *before*
+    it can skip the barrier below (which every other rank is already
+    blocked on, waiting for its turn) — otherwise the whole job hangs
+    instead of failing. The failure is then broadcast so every rank raises
+    together, rather than only the failing rank aborting while the others
+    loop on to a barrier the crashed rank can no longer reach.
+    """
+    if not serial:
+        return build_fn()
+
+    result: T | None = None
+    for i in range(state.num_processes):
+        exc: Exception | None = None
+        if state.process_index == i:
+            try:
+                result = build_fn()
+            except Exception as e:  # noqa: BLE001 - must catch to keep ranks in sync
+                exc = e
+
+        # Reached unconditionally: a failure on rank i must not stop it from
+        # releasing every other rank waiting here for their turn.
+        state.wait_for_everyone()
+
+        if state.num_processes > 1:
+            failed = torch.tensor(int(exc is not None), device=state.device)
+            dist.all_reduce(failed, op=dist.ReduceOp.MAX)
+            if failed.item():
+                if exc is not None:
+                    raise RuntimeError(
+                        f"build_one_at_a_time: rank {i} raised while running build_fn(): {exc!r}"
+                    ) from exc
+                raise RuntimeError(
+                    f"build_one_at_a_time: rank {i} raised while running build_fn(); aborting all ranks."
+                )
+        elif exc is not None:
+            raise exc
+
+    return result
+
+
 _TORCH_DTYPE_MAP: dict[str, torch.dtype] = {
     "bfloat16": torch.bfloat16,
     "float16": torch.float16,

src/post_training/methods/dpo.py

@@ -14,6 +14,7 @@
     build_callbacks,
     build_common_training_kwargs,
     build_model_init_kwargs,
+    build_one_at_a_time,
     build_tokenizer,
     sanitize_generation_config,
 )
@@ -46,7 +47,11 @@ def build_dpo_trainer(config: PostTrainingConfig, run_dir: Path) -> DPOTrainer:
     """
     mc = config.dpo  # method-specific config
 
-    tokenizer = build_tokenizer(config)
+    tokenizer = build_one_at_a_time(
+        PartialState(),
+        lambda: build_tokenizer(config),
+        serial=config.load_model_serially_across_ranks,
+    )
     with PartialState().main_process_first():
         dataset = load_and_mix_datasets(config.data, row_filter=_dpo_row_filter)
 
@@ -59,13 +64,17 @@ def build_dpo_trainer(config: PostTrainingConfig, run_dir: Path) -> DPOTrainer:
         model_init_kwargs=build_model_init_kwargs(config),
     )
 
-    trainer = DPOTrainer(
-        model=config.model.name_or_path,
-        ref_model=mc.ref_model_name_or_path,  # None → TRL creates implicit copy
-        processing_class=tokenizer,
-        train_dataset=dataset,
-        args=dpo_config,
-        callbacks=build_callbacks(config, run_dir),
+    trainer = build_one_at_a_time(
+        PartialState(),
+        lambda: DPOTrainer(
+            model=config.model.name_or_path,
+            ref_model=mc.ref_model_name_or_path,  # None → TRL creates implicit copy
+            processing_class=tokenizer,
+            train_dataset=dataset,
+            args=dpo_config,
+            callbacks=build_callbacks(config, run_dir),
+        ),
+        serial=config.load_model_serially_across_ranks,
     )
     sanitize_generation_config(trainer)
     return trainer

src/post_training/methods/sft.py

@@ -16,6 +16,7 @@
     build_callbacks,
     build_common_training_kwargs,
     build_model_init_kwargs,
+    build_one_at_a_time,
     build_tokenizer,
     sanitize_generation_config,
 )
@@ -59,7 +60,11 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer:
     """
     mc = config.sft  # method-specific config
 
-    tokenizer = build_tokenizer(config)
+    tokenizer = build_one_at_a_time(
+        PartialState(),
+        lambda: build_tokenizer(config),
+        serial=config.load_model_serially_across_ranks,
+    )
 
     # Fail fast if the chat template can't drive `assistant_only_loss=True`.
     # Missing markers silently degrade SFT to full-sequence loss — a 21h run
@@ -108,12 +113,16 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer:
         assistant_only_loss=True,
     )
 
-    trainer = SFTTrainer(
-        model=config.model.name_or_path,
-        processing_class=tokenizer,
-        train_dataset=dataset,
-        args=sft_config,
-        callbacks=build_callbacks(config, run_dir),
+    trainer = build_one_at_a_time(
+        PartialState(),
+        lambda: SFTTrainer(
+            model=config.model.name_or_path,
+            processing_class=tokenizer,
+            train_dataset=dataset,
+            args=sft_config,
+            callbacks=build_callbacks(config, run_dir),
+        ),
+        serial=config.load_model_serially_across_ranks,
     )
 
     sanitize_generation_config(trainer)