From 2bf73ea06973db7070028e4a3f0d2f105ddcbfa7 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 9 Jul 2026 11:11:48 +0200 Subject: [PATCH 1/4] fix: serialize model/tokenizer loading across ranks to avoid HF cache lock contention --- src/post_training/methods/common.py | 24 ++++++++++++++++++++++++ src/post_training/methods/dpo.py | 20 ++++++++++++-------- src/post_training/methods/sft.py | 18 +++++++++++------- 3 files changed, 47 insertions(+), 15 deletions(-) diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index f98a137..8098b74 100644 --- a/src/post_training/methods/common.py +++ b/src/post_training/methods/common.py @@ -9,10 +9,12 @@ import dataclasses import logging import os +from collections.abc import Callable from pathlib import Path from typing import TYPE_CHECKING, Any import torch +from accelerate import PartialState from transformers import AutoTokenizer from post_training.callbacks.inference_checkpoint import InferenceCheckpointCallback @@ -25,6 +27,28 @@ logger = logging.getLogger(__name__) + +def build_one_at_a_time[T](state: PartialState, build_fn: Callable[[], T]) -> 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). + """ + result: T | None = None + for i in range(state.num_processes): + if state.process_index == i: + result = build_fn() + state.wait_for_everyone() + return result + + _TORCH_DTYPE_MAP: dict[str, torch.dtype] = { "bfloat16": torch.bfloat16, "float16": torch.float16, diff --git a/src/post_training/methods/dpo.py b/src/post_training/methods/dpo.py index 14ffb45..bb2c3fb 100644 --- a/src/post_training/methods/dpo.py +++ b/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,7 @@ 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)) with PartialState().main_process_first(): dataset = load_and_mix_datasets(config.data, row_filter=_dpo_row_filter) @@ -59,13 +60,16 @@ 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), + ), ) sanitize_generation_config(trainer) return trainer diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 1607f22..0de8629 100644 --- a/src/post_training/methods/sft.py +++ b/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,7 @@ 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)) # 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 +109,15 @@ 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), + ), ) sanitize_generation_config(trainer) From 7edd27ab6146090563c8d75ce0b7724fb7081616 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 9 Jul 2026 12:43:58 +0200 Subject: [PATCH 2/4] fix: handle None result in build_one_at_a_time --- src/post_training/methods/common.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index 8098b74..c4aa1b9 100644 --- a/src/post_training/methods/common.py +++ b/src/post_training/methods/common.py @@ -46,6 +46,7 @@ def build_one_at_a_time[T](state: PartialState, build_fn: Callable[[], T]) -> T: if state.process_index == i: result = build_fn() state.wait_for_everyone() + assert result is not None, "build_one_at_a_time: result is None." return result From 1ce8bcc48aa106db6082d84b60c955893e3a279c Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Mon, 13 Jul 2026 12:30:59 +0200 Subject: [PATCH 3/4] feat: make serialization of model load optional --- configs/trl/dpo.yaml | 1 + configs/trl/sft.yaml | 1 + src/post_training/config.py | 1 + src/post_training/methods/common.py | 11 ++++++++++- src/post_training/methods/dpo.py | 7 ++++++- src/post_training/methods/sft.py | 7 ++++++- 6 files changed, 25 insertions(+), 3 deletions(-) diff --git a/configs/trl/dpo.yaml b/configs/trl/dpo.yaml index 10c6433..1df1102 100644 --- a/configs/trl/dpo.yaml +++ b/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: diff --git a/configs/trl/sft.yaml b/configs/trl/sft.yaml index 90566fc..0f942f4 100644 --- a/configs/trl/sft.yaml +++ b/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: diff --git a/src/post_training/config.py b/src/post_training/config.py index f98bc63..ea0be0e 100644 --- a/src/post_training/config.py +++ b/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) diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index c4aa1b9..0dc9954 100644 --- a/src/post_training/methods/common.py +++ b/src/post_training/methods/common.py @@ -28,7 +28,9 @@ logger = logging.getLogger(__name__) -def build_one_at_a_time[T](state: PartialState, build_fn: Callable[[], T]) -> T: +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 @@ -40,7 +42,14 @@ def build_one_at_a_time[T](state: PartialState, build_fn: Callable[[], T]) -> T: 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 not serial: + return build_fn() + result: T | None = None for i in range(state.num_processes): if state.process_index == i: diff --git a/src/post_training/methods/dpo.py b/src/post_training/methods/dpo.py index bb2c3fb..85c2877 100644 --- a/src/post_training/methods/dpo.py +++ b/src/post_training/methods/dpo.py @@ -47,7 +47,11 @@ def build_dpo_trainer(config: PostTrainingConfig, run_dir: Path) -> DPOTrainer: """ mc = config.dpo # method-specific config - tokenizer = build_one_at_a_time(PartialState(), lambda: 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) @@ -70,6 +74,7 @@ def build_dpo_trainer(config: PostTrainingConfig, run_dir: Path) -> DPOTrainer: args=dpo_config, callbacks=build_callbacks(config, run_dir), ), + serial=config.load_model_serially_across_ranks, ) sanitize_generation_config(trainer) return trainer diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 0de8629..ca7fb89 100644 --- a/src/post_training/methods/sft.py +++ b/src/post_training/methods/sft.py @@ -60,7 +60,11 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer: """ mc = config.sft # method-specific config - tokenizer = build_one_at_a_time(PartialState(), lambda: 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 @@ -118,6 +122,7 @@ def build_sft_trainer(config: PostTrainingConfig, run_dir: Path) -> SFTTrainer: args=sft_config, callbacks=build_callbacks(config, run_dir), ), + serial=config.load_model_serially_across_ranks, ) sanitize_generation_config(trainer) From d0334f08357ca0684dea115f9dcacae604a453e4 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Mon, 13 Jul 2026 13:47:43 +0200 Subject: [PATCH 4/4] fix: prevent deadlock in build_one_at_a_time on build_fn failure If build_fn() raised on the active rank, it skipped wait_for_everyone(), leaving every other rank blocked on that barrier forever. Now the failure is caught, the barrier is still reached, and the failure is broadcast so all ranks raise together instead of hanging. --- src/post_training/methods/common.py | 33 +++++++++++++++++++++++++++-- 1 file changed, 31 insertions(+), 2 deletions(-) diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index 0dc9954..0062d3b 100644 --- a/src/post_training/methods/common.py +++ b/src/post_training/methods/common.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any import torch +import torch.distributed as dist from accelerate import PartialState from transformers import AutoTokenizer @@ -46,16 +47,44 @@ def build_one_at_a_time[T]( 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: - result = build_fn() + 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() - assert result is not None, "build_one_at_a_time: result is None." + + 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