fix: shared memory

README.md

@@ -101,37 +101,37 @@ 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
 ```
 
 #### 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
 

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,