Add data-pruning scaling-law experiment pipeline (Telugu) - #1
Open
vishnup22 wants to merge 11 commits into
Open
Conversation
Implements the core research question: does pruning 30-50% of low-resource pretraining data with a small reference model preserve baseline BPB? - evaluation/perplexity.py: score_lines() for per-line (not batch-averaged) cross-entropy scoring, verified against eval_perplexity and padding-safe. - pruning/score.py: streams a train split through a reference checkpoint and writes per-line loss scores. - pruning/make_splits.py: partitions a scored split into easy/hard/mid (Pareto core-set)/random 50% subsets; verified against synthetic data. - models/gpt2/train.py: --variant/--eval_strategy/--eval_steps/--seed flags to train on pruned splits with step-level eval logging (eval_history: tokens_seen/eval_loss/bpb) for the tokens-vs-BPB curve; fully backward compatible when --variant is omitted. - pruning/evaluate_variants.py: BPB/perplexity + WikiANN NER F1 for every variant plus the existing 100% baseline (re-evaluated from the HF Hub checkpoint for a fair final-number comparison). - analysis/plot_pruning_scaling.py + summarize_pruning.py: the deliverable tokens-vs-BPB plot and results table. - scripts/run_pruning_experiment.sh, configs/experiments/telugu-pruning-scaling.yaml, docs/pruning_experiment.md: SLURM launcher, experiment manifest, methodology writeup. - matplotlib added as a dependency.
The launcher was configured for 1 GPU per variant, which (combined with compute_total_steps scaling inversely with world_size) would have meant ~2x the baseline's step count per variant despite training on half the data -- ~5.6 days/variant, ~22-24 days sequential, over the 14-day budget. Switch to accelerate launch --num_processes 4 (matching train_gpt.sh and the telugu_seed2 baseline), which brings each variant back to roughly 1.4 days (half the baseline's 67.1h on the same 4-GPU config), ~5.6 days for all 4 sequentially. Bumped --gres to gpu:4 and reduced --time from 14 to 8 days accordingly.
…r.model T5Tokenizer.from_pretrained() resolves vocab_files_names['vocab_file'], which is hardcoded to 'spiece.model' -- but the Hub repo's raw SentencePiece file is named 'tokenizer.model' (matching how models.gpt2.train saves it), so auto-resolution silently returned None and crashed inside sentencepiece with a non-obvious 'TypeError: not a string'. Download the known filename explicitly via hf_hub_download, same pattern train.py's load_tokenizer already uses for local checkpoints. Verified against the real pulipakav-1/dravidian-gpt2-telugu repo: loads, vocab_size 32000, correct pad_token_id, encodes/decodes correctly.
--time was set to 8-00:00:00 but the cluster's actual cap is 7 days -- tight against the ~5.6-7 day estimate for the full pipeline, with no margin for a mid-run timeout. train_one() previously only skipped a variant if its result JSON already existed, so a timeout mid-training would restart that variant from scratch on resubmission, silently wasting however many days it had already run. Now it resumes from the latest checkpoint (trainer.train(resume_from_checkpoint=True)) when output_dir has checkpoints but no finished result JSON, so the pipeline is resumable at both the per-variant and mid-variant level -- just re-sbatch the same script if it times out.
conda's activation hook references unset variables internally, so -u (nounset) makes 'conda activate telugu_llm' abort the script. -eo pipefail still catches command failures and pipe failures; just not unset vars.
…odel loads Passing dtype=/torch_dtype= directly into .from_pretrained() leaks that kwarg into AutoConfig/PretrainedConfig's unused_kwargs on this transformers version, gets stored as a stray non-JSON-serializable attribute on the config, and crashes the first time the config is repr'd/logged (which from_dict() does unconditionally when return_unused_kwargs is set) -- regardless of which model actually gets loaded. Fixed all four call sites the same way: load without a dtype kwarg, then .to(dtype=...) after. Verified against the real HF Hub checkpoint (pulipakav-1/dravidian-gpt2-telugu) -- previously reproduced the exact traceback reported from the cluster, now loads cleanly. - evaluation/run_eval.py::load_model_and_tokenizer (the one that crashed) - pruning/evaluate_variants.py::load_variant (same pattern, not yet exercised) - evaluation/downstream.py::_load_seq_clf_model, run_wikiann_ner (same pattern with the old torch_dtype= name, would have hit this next in pruning.evaluate_variants' downstream eval step)
…ward pass) Observed on the cluster: 198/~44,700 chunks in 6min at batch_size=8, which projects to ~22.6h for the full scoring pass -- too much of the 7-day budget for a step that's pure forward-pass inference. score_lines() is already @torch.no_grad(), so there's no gradient/optimizer memory pressure to justify a small batch size; bump it way up. Also bumped READ_CHUNK_LINES 512 -> 2048 to reduce per-chunk Python loop overhead.
…when tokenizer.model isn't already present
…tegy Newer transformers versions renamed evaluation_strategy -> eval_strategy. Claude-Session: https://claude.ai/code/session_01F3trZ6zLPewuWpFRUPoisx
…lready done Steps 1 (scoring, ~7.8h for te_train) and 2 (pruned split building) had no idempotency check, so resubmitting the script after a crash further into the pipeline (e.g. the just-fixed evaluation_strategy TypeError) would silently redo the expensive scoring pass for no reason. Guard both on their output files existing, matching the tokenizer step's existing skip-if-present behavior and the resumability the script's header comment already promised. Claude-Session: https://claude.ai/code/session_01F3trZ6zLPewuWpFRUPoisx
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Research question
Does pruning 30-50% of low-resource pretraining data using a small reference model preserve (or beat) 100%-baseline validation BPB — i.e. does it break standard power-law scaling? Motivated by Ari Morcos' NeurIPS finding (easy examples help when data is scarce, hard examples help at scale) and Datology's small-reference-model data curation work.
Full methodology:
docs/pruning_experiment.md.What this adds
pruning/score.py— streams a train split through a reference checkpoint (default:pulipakav-1/dravidian-gpt2-telugu) and writes per-line cross-entropy loss.pruning/make_splits.py— partitions scored lines into 50%easy/hard/mid(40th-90th percentile Pareto core-set) /randomsubsets. Val/test splits are untouched so every variant is judged on identical held-out data.evaluation/perplexity.py: score_lines()— new per-line (not batch-averaged) loss function thatscore.pydepends on. Numerically verified against the existingeval_perplexityaggregate and confirmed padding-safe.models/gpt2/train.py— new--variant/--eval_strategy/--eval_steps/--seedflags to train on a pruned split with step-level eval logging (eval_history: tokens_seen / eval_loss / bpb) for the tokens-vs-BPB curve. Omitting--variantreproduces prior behavior/filenames exactly (backward compatible).pruning/evaluate_variants.py— BPB/perplexity + WikiANN NER F1 for every trained variant plus the existing 100% baseline (re-evaluated fresh from the HF Hub checkpoint for a fair final-number comparison; its training curve stays the existing coarse one — see docs for why).analysis/plot_pruning_scaling.py+summarize_pruning.py— the deliverable tokens-vs-BPB plot and a results table.scripts/run_pruning_experiment.sh(SLURM launcher),configs/experiments/telugu-pruning-scaling.yaml(experiment manifest), README pointers,matplotlibadded as a dependency.Verification
No GPU/corpus available in the dev environment, so:
make_splits.pyrun end-to-end on synthetic scores/text: clean 50/50 easy/hard partition, correct percentile thresholds formid, no cross-split overlap, blank lines correctly excluded.score_lineschecked against a real (tiny) GPT-2 forward pass: matcheseval_perplexity's aggregate loss exactly on single-example batches, and confirmed padding never leaks into another line's score.build_eval_historyunit-checked for correcttokens_seen/bpbcomputation and non-eval log entry filtering.Full-scale scoring + training (4 variants + baseline eval) needs the SLURM cluster — run via
scripts/run_pruning_experiment.sh, documented indocs/pruning_experiment.md.