From de27b7208a2ccb3a97c9e48d3b79c03f32c5eaa4 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sun, 8 Feb 2026 09:48:24 -0800 Subject: [PATCH 1/2] =?UTF-8?q?fix:=20resolve=20safetensors=20shared=20ten?= =?UTF-8?q?sor=20error=20and=20migrate=20to=20accelerate=20SplitHeadWrappe?= =?UTF-8?q?r=20wraps=20RobertaForMaskedLM=20which=20has=20tied=20weights?= =?UTF-8?q?=20(word=5Fembeddings.weight=20=E2=86=94=20lm=5Fhead.decoder.we?= =?UTF-8?q?ight).=20Safetensors=20rejects=20shared=20tensors=20during=20ch?= =?UTF-8?q?eckpoint=20saving.=20Switch=20to=20torch.save=20via=20save=5Fsa?= =?UTF-8?q?fetensors=3DFalse.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace torch.distributed with accelerate.PartialState for world-size detection, and torchrun with accelerate launch for cleaner distributed training management.ix: resolve safetensors shared tensor error and migrate to accelerate SplitHeadWrapper wraps RobertaForMaskedLM which has tied weights (word_embeddings.weight ↔ lm_head.decoder.weight). Safetensors rejects shared tensors during checkpoint saving. Switch to torch.save via save_safetensors=False. Replace torch.distributed with accelerate.PartialState for world-size detection, and torchrun with accelerate launch for cleaner distributed training management. --- modeling/pretrain.py | 17 ++++++++--------- run_pretrain.sh | 7 ------- 2 files changed, 8 insertions(+), 16 deletions(-) delete mode 100755 run_pretrain.sh diff --git a/modeling/pretrain.py b/modeling/pretrain.py index aad05ef..1a912fb 100644 --- a/modeling/pretrain.py +++ b/modeling/pretrain.py @@ -5,8 +5,7 @@ from collections import defaultdict from functools import partial -import torch -import torch.distributed as dist +from accelerate import PartialState from datasets import Dataset, Features, Value, load_dataset from transformers import ( DataCollatorForLanguageModeling, @@ -23,10 +22,8 @@ def _get_world_size() -> int: - """Return DDP world size if distributed is initialized, else GPU count (min 1).""" - if dist.is_initialized(): - return dist.get_world_size() - return max(torch.cuda.device_count(), 1) + """Return world size via Accelerate (handles both distributed and single-process).""" + return PartialState().num_processes def compute_function_id(code: str) -> int: @@ -231,9 +228,9 @@ def main( # Cap num_proc to a sane default (see modeling.common.default_num_proc()). num_proc = min(num_proc, default_num_proc()) - # When launched with torchrun, each rank would otherwise spawn num_proc - # workers, quickly oversubscribing CPUs (e.g., 4 ranks × 80 workers = 320). - world_size = int(os.environ.get("WORLD_SIZE", "1")) + # Each rank would otherwise spawn num_proc workers, quickly oversubscribing + # CPUs (e.g., 4 ranks × 80 workers = 320). + world_size = _get_world_size() if world_size > 1: num_proc = max(1, num_proc // world_size) if num_proc > 1: @@ -330,6 +327,7 @@ def main( save_total_limit=3, load_best_model_at_end=True, dataloader_num_workers=max(1, (os.cpu_count() or 1) // _get_world_size()), + save_safetensors=False, # SplitHeadWrapper has tied weights from RobertaForMaskedLM ) trainer = ContrastiveTrainer( @@ -338,6 +336,7 @@ def main( train_dataset=train_dataset, eval_dataset=eval_dataset, data_collator=collator_fn, + processing_class=tokenizer, alpha=alpha, temperature=temperature, contra_mode=contra_mode, diff --git a/run_pretrain.sh b/run_pretrain.sh deleted file mode 100755 index d67ff65..0000000 --- a/run_pretrain.sh +++ /dev/null @@ -1,7 +0,0 @@ -#!/bin/bash - -# export CUDA_VISIBLE_DEVICES=4,5,6,7 -export WANDB_PROJECT="InvPT" - -# Use torchrun for DDP; uses all visible GPUs (control with CUDA_VISIBLE_DEVICES). -torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml From 7af0c522d6301a92dbed2fb27df97682757c24c9 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sun, 8 Feb 2026 09:53:03 -0800 Subject: [PATCH 2/2] doc: update to accelerate --- README.md | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 0c18e03..c13f1f4 100755 --- a/README.md +++ b/README.md @@ -101,23 +101,23 @@ The CLI entry point is `modeling/cli.py`, which provides two subcommands: - **`run`** -- load a YAML experiment config (recommended; all parameters come from the config file to ensure full reproducibility) - **`pretrain`** -- pass all parameters directly as CLI options -Pre-training uses PyTorch DistributedDataParallel (DDP) via `torchrun`, which ships with PyTorch itself (no extra dependencies). The same training code runs on both multi-GPU and single-GPU nodes -- the HuggingFace `Trainer` auto-detects the distributed environment set up by `torchrun` and enables or disables DDP accordingly. +Pre-training uses PyTorch DistributedDataParallel (DDP) via [HuggingFace Accelerate](https://huggingface.co/docs/accelerate). The same training code runs on both multi-GPU and single-GPU nodes -- the HuggingFace `Trainer` auto-detects the distributed environment and enables or disables DDP accordingly. #### Multi-GPU node -Use `torchrun` to spawn one process per GPU. `--nproc_per_node=gpu` automatically uses all visible GPUs: +Use `accelerate launch --multi_gpu` to spawn one process per GPU. It automatically uses all visible GPUs: ```sh -torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml +accelerate launch --multi_gpu modeling/cli.py run experiments/base.yaml ``` To select specific GPUs, set `CUDA_VISIBLE_DEVICES`: ```sh -CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=gpu modeling/cli.py run experiments/base.yaml +CUDA_VISIBLE_DEVICES=0,1 accelerate launch --multi_gpu modeling/cli.py run experiments/base.yaml ``` -There is also a convenience script that launches `torchrun` on all visible GPUs: +There is also a convenience script that launches on all visible GPUs: ```sh ./run_pretrain.sh @@ -125,13 +125,13 @@ There is also a convenience script that launches `torchrun` on all visible GPUs: #### Single-GPU node -On a single-GPU machine, run with plain `python` -- no `torchrun` needed: +On a single-GPU machine, run with plain `python` -- no `accelerate launch` needed: ```sh python modeling/cli.py run experiments/base.yaml ``` -`torchrun --nproc_per_node=1` also works if you prefer a uniform launch command across environments. +`accelerate launch` (without `--multi_gpu`) also works if you prefer a uniform launch command across environments. #### CLI examples @@ -215,7 +215,7 @@ The `pretrain` subcommand accepts all training parameters directly as CLI option | `--contra-mode` | `info_nce` | Contrastive loss mode: `info_nce`, `supcon`, or `grouped` | | `--max-num-augs` | `6` | Max augmentations per anchor group (`grouped` mode only) | -Note: dataset preprocessing uses HuggingFace Datasets multiprocessing. When running with `torchrun` (multi-GPU), `--num-proc` is automatically scaled down per-rank to avoid CPU oversubscription, and `TOKENIZERS_PARALLELISM` is disabled when using multiple workers. +Note: dataset preprocessing uses HuggingFace Datasets multiprocessing. When running multi-GPU, `--num-proc` is automatically scaled down per-rank to avoid CPU oversubscription, and `TOKENIZERS_PARALLELISM` is disabled when using multiple workers. #### Contrastive Loss Modes