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 f98a137..0062d3b 100644 --- a/src/post_training/methods/common.py +++ b/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, diff --git a/src/post_training/methods/dpo.py b/src/post_training/methods/dpo.py index 14ffb45..85c2877 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,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 diff --git a/src/post_training/methods/sft.py b/src/post_training/methods/sft.py index 1607f22..ca7fb89 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,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)