update max num proc

README.md

@@ -210,11 +210,13 @@ The `pretrain` subcommand accepts all training parameters directly as CLI option
 | `--sample-rate` | `1.0` | Fraction of dataset to use (for quick experiments) |
 | `--seed` | `0` | Random seed |
 | `--run-name` | `InvariantBERT` | W&B run name and output directory name |
-| `--num-proc` | `80` | Number of processes for dataset tokenization |
+| `--num-proc` | (all CPU cores) | Parallel workers for dataset preprocessing |
 | `--resume / --no-resume` | `False` | Resume training from the latest checkpoint |
 | `--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.
+
 #### Contrastive Loss Modes
 
 The `--contra-mode` option selects the contrastive loss function:

experiments/grouped/codebert.yaml

@@ -5,9 +5,9 @@
 dataset_path: "data/aug_csn.jsonl"
 model_name: "microsoft/codebert-base"
 
-batch_size: 64
+batch_size: 32
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 8
 learning_rate: 2.0e-5
 
 seed: 0

experiments/grouped/contrabert_c.yaml

@@ -6,9 +6,9 @@ dataset_path: "data/aug_csn.jsonl"
 model_name: "./saved_models/ContraBERT_C"
 tokenizer_name: "microsoft/codebert-base"
 
-batch_size: 64
+batch_size: 32
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 8
 learning_rate: 2.0e-5
 
 seed: 0

experiments/grouped/contrabert_g.yaml

@@ -6,9 +6,9 @@ dataset_path: "data/aug_csn.jsonl"
 model_name: "./saved_models/ContraBERT_G"
 tokenizer_name: "microsoft/graphcodebert-base"
 
-batch_size: 64
+batch_size: 32
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 8
 learning_rate: 2.0e-5
 
 seed: 0

experiments/grouped/graphcodebert.yaml

@@ -5,9 +5,9 @@
 dataset_path: "data/aug_csn.jsonl"
 model_name: "microsoft/graphcodebert-base"
 
-batch_size: 64
+batch_size: 32
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 8
 learning_rate: 2.0e-5
 
 seed: 0

experiments/supcon/codebert.yaml

@@ -5,9 +5,9 @@
 dataset_path: "data/aug_csn.jsonl"
 model_name: "microsoft/codebert-base"
 
-batch_size: 64
+batch_size: 256
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 1
 learning_rate: 2.0e-5
 
 seed: 0

experiments/supcon/contrabert_c.yaml

@@ -6,9 +6,9 @@ dataset_path: "data/aug_csn.jsonl"
 model_name: "./saved_models/ContraBERT_C"
 tokenizer_name: "microsoft/codebert-base"
 
-batch_size: 64
+batch_size: 256
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 1
 learning_rate: 2.0e-5
 
 seed: 0

experiments/supcon/contrabert_g.yaml

@@ -6,9 +6,9 @@ dataset_path: "data/aug_csn.jsonl"
 model_name: "./saved_models/ContraBERT_G"
 tokenizer_name: "microsoft/graphcodebert-base"
 
-batch_size: 64
+batch_size: 256
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 1
 learning_rate: 2.0e-5
 
 seed: 0

experiments/supcon/graphcodebert.yaml

@@ -5,9 +5,9 @@
 dataset_path: "data/aug_csn.jsonl"
 model_name: "microsoft/graphcodebert-base"
 
-batch_size: 64
+batch_size: 256 # 64
 num_epochs: 3
-gradient_accumulation_steps: 4
+gradient_accumulation_steps: 1 # 4
 learning_rate: 2.0e-5
 
 seed: 0

modeling/cli.py

@@ -71,7 +71,10 @@ def pretrain(
         float, typer.Option(help="Fraction of dataset to sample.")
     ] = 1.0,
     num_proc: Annotated[
-        int, typer.Option(help="Number of dataloader workers.")
+        int,
+        typer.Option(
+            help="Parallel workers for dataset preprocessing (datasets filter/map)."
+        ),
     ] = default_num_proc(),
     resume: Annotated[
         bool, typer.Option(help="Resume from latest checkpoint.")

modeling/common.py

@@ -4,11 +4,15 @@
 import numpy as np
 import torch
 
-MAX_NUM_PROC = 80
-
 
 def default_num_proc() -> int:
-    """Return the default number of parallel workers, capped at available CPUs."""
+    """Return the default number of parallel workers.
+
+    HuggingFace Datasets preprocessing (e.g., ``Dataset.map(num_proc=...)``) and
+    tokenizers can both parallelize. Using very high ``num_proc`` values can
+    oversubscribe CPU and/or hammer the datasets cache on disk, often making
+    preprocessing *slower*.
+    """
     return os.cpu_count() or 1
 
 

modeling/pretrain.py

@@ -38,14 +38,12 @@ def compute_function_id(code: str) -> int:
 def tokenize(tokenizer, example, max_seq_length=256):
     code_inputs = tokenizer(
         example["code"],
-        padding="max_length",
         truncation=True,
         max_length=max_seq_length,
         return_special_tokens_mask=True,
     )
     aug_inputs = tokenizer(
         example["transformed"],
-        padding="max_length",
         truncation=True,
         max_length=max_seq_length,
         return_special_tokens_mask=True,
@@ -90,8 +88,8 @@ def regroup_dataset(dataset, max_num_augs: int = 6) -> Dataset:
         }
     )
 
-    for i in range(len(dataset)):
-        row = dataset[i]
+    # Iterating a HF Dataset is much faster than random indexing (dataset[i]).
+    for row in dataset:
         fid = compute_function_id(row["code"])
         g = groups[fid]
         if g["code"] is None:
@@ -231,8 +229,16 @@ def main(
 ):
     set_seed(seed)
 
-    # Cap num_proc to available CPU cores to avoid broken-pipe errors.
+    # 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"))
+    if world_size > 1:
+        num_proc = max(1, num_proc // world_size)
+    if num_proc > 1:
+        # Avoid oversubscription (processes × tokenizer threads).
+        os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
     tokenizer_name = tokenizer_name or model_name
     tokenizer = RobertaTokenizerFast.from_pretrained(tokenizer_name)
@@ -261,7 +267,11 @@ def main(
         }
     )
     dataset = load_dataset("json", data_files=dataset_path, features=features)["train"]
-    dataset = dataset.filter(lambda x: x["transformed"] is not None)
+    dataset = dataset.filter(
+        lambda transformed: transformed is not None,
+        input_columns=["transformed"],
+        num_proc=num_proc,
+    )
 
     if sample_rate < 1.0:
         dataset = dataset.shuffle(seed=seed).select(
@@ -284,6 +294,7 @@ def main(
             ),
             batched=True,
             num_proc=num_proc,
+            remove_columns=grouped_dataset.column_names,
         ).shuffle(seed=seed)
 
         collator_fn = partial(grouped_contra_data_collator, mlm_collator, max_num_augs)
@@ -292,6 +303,7 @@ def main(
             partial(tokenize, tokenizer, max_seq_length=max_seq_length),
             batched=True,
             num_proc=num_proc,
+            remove_columns=dataset.column_names,
         ).shuffle(seed=seed)
 
         collator_fn = partial(contra_data_collator, mlm_collator)

tests/test_common.py

@@ -0,0 +1,13 @@
+import os
+
+from modeling.common import default_num_proc
+
+
+def test_default_num_proc_returns_cpu_count(monkeypatch) -> None:
+    monkeypatch.setattr(os, "cpu_count", lambda: 123)
+    assert default_num_proc() == 123
+
+
+def test_default_num_proc_is_at_least_one(monkeypatch) -> None:
+    monkeypatch.setattr(os, "cpu_count", lambda: None)
+    assert default_num_proc() == 1