From 07899711d3c0b4c50f468df828f6ab1694b6b816 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sat, 7 Feb 2026 18:18:17 -0800 Subject: [PATCH 1/3] update max num proc --- README.md | 4 +++- modeling/cli.py | 5 ++++- modeling/common.py | 10 +++++++--- modeling/pretrain.py | 24 ++++++++++++++++++------ tests/test_common.py | 13 +++++++++++++ 5 files changed, 45 insertions(+), 11 deletions(-) create mode 100644 tests/test_common.py diff --git a/README.md b/README.md index cc2eb24..0c18e03 100755 --- a/README.md +++ b/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: diff --git a/modeling/cli.py b/modeling/cli.py index c12165c..28cef55 100644 --- a/modeling/cli.py +++ b/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.") diff --git a/modeling/common.py b/modeling/common.py index ed9cbe3..3f31259 100644 --- a/modeling/common.py +++ b/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 diff --git a/modeling/pretrain.py b/modeling/pretrain.py index 7165671..aad05ef 100644 --- a/modeling/pretrain.py +++ b/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) diff --git a/tests/test_common.py b/tests/test_common.py new file mode 100644 index 0000000..854d589 --- /dev/null +++ b/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 From c9966f9204546669f3daa4c224ffa60dfc406045 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sun, 8 Feb 2026 11:46:05 +0800 Subject: [PATCH 2/3] chore: use 32 batch size for grouped --- experiments/grouped/codebert.yaml | 4 ++-- experiments/grouped/contrabert_c.yaml | 4 ++-- experiments/grouped/contrabert_g.yaml | 4 ++-- experiments/grouped/graphcodebert.yaml | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/experiments/grouped/codebert.yaml b/experiments/grouped/codebert.yaml index 353590e..b414305 100644 --- a/experiments/grouped/codebert.yaml +++ b/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 diff --git a/experiments/grouped/contrabert_c.yaml b/experiments/grouped/contrabert_c.yaml index 93432a4..a7456e4 100644 --- a/experiments/grouped/contrabert_c.yaml +++ b/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 diff --git a/experiments/grouped/contrabert_g.yaml b/experiments/grouped/contrabert_g.yaml index b30a708..d9ae460 100644 --- a/experiments/grouped/contrabert_g.yaml +++ b/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 diff --git a/experiments/grouped/graphcodebert.yaml b/experiments/grouped/graphcodebert.yaml index 7104fbf..da9ff73 100644 --- a/experiments/grouped/graphcodebert.yaml +++ b/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 From 4c0ed401ac60fb4e317cc4aec95251bdf34156d3 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sun, 8 Feb 2026 14:28:30 +0800 Subject: [PATCH 3/3] chore: use 256 batchsize for supcon --- experiments/supcon/codebert.yaml | 4 ++-- experiments/supcon/contrabert_c.yaml | 4 ++-- experiments/supcon/contrabert_g.yaml | 4 ++-- experiments/supcon/graphcodebert.yaml | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/experiments/supcon/codebert.yaml b/experiments/supcon/codebert.yaml index 6831157..fb70d1a 100644 --- a/experiments/supcon/codebert.yaml +++ b/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 diff --git a/experiments/supcon/contrabert_c.yaml b/experiments/supcon/contrabert_c.yaml index 0694e4c..039e83f 100644 --- a/experiments/supcon/contrabert_c.yaml +++ b/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 diff --git a/experiments/supcon/contrabert_g.yaml b/experiments/supcon/contrabert_g.yaml index 5bc03e7..eb262b6 100644 --- a/experiments/supcon/contrabert_g.yaml +++ b/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 diff --git a/experiments/supcon/graphcodebert.yaml b/experiments/supcon/graphcodebert.yaml index 6360c18..8d8c835 100644 --- a/experiments/supcon/graphcodebert.yaml +++ b/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