fix: serialize model/tokenizer loading across ranks to avoid HF cache contention#45
fix: serialize model/tokenizer loading across ranks to avoid HF cache contention#45Neonkraft wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
Serializes Hugging Face Hub cache resolution across distributed ranks to prevent intermittent from_pretrained cache/lock contention errors observed on Lustre-backed shared filesystems (e.g., LUMI), by forcing rank-ordered construction of tokenizers and trainers.
Changes:
- Add
build_one_at_a_time()utility that runs a build function one rank at a time usingaccelerate.PartialStatebarriers. - Wrap tokenizer construction in
sft.pyanddpo.pywithbuild_one_at_a_time(...)to avoid concurrent HF cache resolution. - Wrap TRL trainer construction in
sft.pyanddpo.pywithbuild_one_at_a_time(...)to serialize model loading paths invoked by TRL.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| src/post_training/methods/common.py | Adds rank-ordered barrier helper to serialize cache-sensitive build steps. |
| src/post_training/methods/sft.py | Uses the helper to serialize tokenizer and SFTTrainer construction across ranks. |
| src/post_training/methods/dpo.py | Uses the helper to serialize tokenizer and DPOTrainer construction across ranks. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 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 |
|
I have one concern. The loading of model/ dataset is now serialized across rank but does that not mean the the time to startup is going to blow as you increase number of ranks to train on? I'm just thinking out if this is the best solution to handle this behavior. I can imagine doing a local stage for each node but I'm not sure if thats an effort worth spending- Basically one rank may copy the data acorss all the nodes so that theres no race and filelock occurring. Or other one that might be bit more complicated would be one rank loading it and broadcasting it to the others (FSDP style). But if the time is not a concern, then this is a fair and valid solution, and you can merge the PR. P.S. I did not run on LUMI as the compute nodes are blocked as of now! |
Only the instantiation of the models and tokenizers are serialized. It might cost roughly half a minute per rank. The solution isn't elegant, but I'd stick to it for now. We can optimize this in future if it becomes an impactful bottleneck. |
Summary
On multi-GPU/multi-rank runs, every rank independently called
from_pretrainedfor the tokenizer and model, each resolving the same HuggingFace Hub repo throughhuggingface_hub'sfilelock-guarded cache-verification path at the same time. On the LUMI cluster (Lustre-backed shared filesystem), this N-way concurrent lock contention intermittently produced spuriousOSError: <repo> does not appear to have a file named ...errors on shard files that were, in fact, present and had been fully cached hours earlier — the failure occurs during cache resolution/locking, not during actual weight loading (another rank in the same job would already be mid-way through loading weights from the same file).This PR adds
build_one_at_a_time()(src/post_training/methods/common.py) and wires it around the tokenizer and trainer (model) construction in bothsft.pyanddpo.py. It forces ranks through the HF cache resolution one at a time via a rank-orderedaccelerate.PartialStatebarrier, so no two ranks ever contend for the same cache lock simultaneously. The barrier synchronizes over the NCCL/Gloo interconnect, not the filesystem, so its correctness doesn't depend on the same Lustre locking behavior that's suspected to be unreliable under concurrent access.Type of change
Validation
Reproduced on the LUMI cluster by submitting the same training job 3 times without the fix and 3 times with it:
OSError: ... does not appear to have a file named model-0000X-of-0000Y.safetensorsrace during multi-rank model loading.