{"id":31908799,"url":"https://github.com/allenai/staged-training","last_synced_at":"2025-10-13T15:56:42.380Z","repository":{"id":42967996,"uuid":"465620077","full_name":"allenai/staged-training","owner":"allenai","description":"Staged Training for Transformer Language Models","archived":false,"fork":false,"pushed_at":"2022-03-31T00:42:26.000Z","size":331,"stargazers_count":32,"open_issues_count":2,"forks_count":2,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-09-05T01:51:37.589Z","etag":null,"topics":["deep-learning","language-model","nlp","pytorch","transformers"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/allenai.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2022-03-03T07:51:28.000Z","updated_at":"2025-01-28T04:02:36.000Z","dependencies_parsed_at":"2022-09-08T03:32:03.761Z","dependency_job_id":null,"html_url":"https://github.com/allenai/staged-training","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/allenai/staged-training","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2Fstaged-training","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2Fstaged-training/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2Fstaged-training/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2Fstaged-training/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/allenai","download_url":"https://codeload.github.com/allenai/staged-training/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/allenai%2Fstaged-training/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279015939,"owners_count":26085777,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-13T02:00:06.723Z","response_time":61,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","language-model","nlp","pytorch","transformers"],"created_at":"2025-10-13T15:56:41.045Z","updated_at":"2025-10-13T15:56:42.374Z","avatar_url":"https://github.com/allenai.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# staged-training\n\nIn our paper [**Staged Training for Transformer Language Models**](https://arxiv.org/abs/2203.06211), we propose a staged training setup that begins with a small model and incrementally increases the amount of compute used for training by applying a \"growth operator\" to increase the model depth and width. By initializing each stage with the output of the previous one, the training process effectively re-uses the compute from prior stages and becomes more efficient. \n\nWe release the reproducible code for the growth operator and evaluation scripts here.\n\n## Setup\n\nThe scripts in this repository require Python 3.7 or newer.\nOnce you have a suitable Python environment, first install PyTorch v1.9.0 according the [official instructions](https://pytorch.org/get-started/previous-versions/#v190). Then run\n```\npip install -r requirements.txt\n```\n\n## Growth Operator\n\nOur growth operators (width/depth) each take as input the entire training state (including model parameters, optimizer state, learning rate schedule, etc.) and output a new training state from which training continues.\n\nPlease see the `scripts/cheatsheet.txt` for more examples on how to use the corresponding scripts. \n\nFor example, you can apply the width operator with:\n```\nCUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/gpt_pretrain.py \\\n  --save_prefix final_gpt2_large_div2_width_check_bs512_lr0.0020_warmup3k_seqlen1024_debug \\\n  --gpu_count -1 \\\n  --model gpt2  \\\n  --tokenizer gpt2 \\\n  --batch_size 4 \\\n  --grad_accum 32  \\\n  --lr 0.002006911598778545  \\\n  --warmup_steps 3000 \\  \\\n  --train_steps 250000  \\\n  --val_every 50  \\\n  --val_batches 50 \\\n  --fp16 \\\n  --seqlen 1024 \\\n  --log_rate 10 \\\n  --num_workers 4 \\\n  --size GPT2_large_div2_width \\\n  --random \\\n  --resume final_runs/final_gpt2_large_div2_width_check_bs512_lr0.0021_warmup3k_seqlen1024_debug/checkpoint-xxx.ckpt \\\n  --doubling weights\n```\n\nOr the depth operator with:\n```\nCUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/gpt_pretrain.py \\\n  --save_prefix final_gpt2_large_div2_depthx2_check_bs512_lr0.0020_warmup3k_seqlen1024_debug \\\n  --gpu_count -1 \\\n  --model gpt2  \\\n  --tokenizer gpt2 \\\n  --batch_size 4 \\\n  --grad_accum 32 \\\n  --lr 0.002006911598778545 \\\n  --warmup_steps 3000 \\\n  --train_steps 250000 \\\n  --val_every 50 \\\n  --val_batches 50 \\\n  --fp16 \\\n  --seqlen 1024 \\\n  --log_rate 10 \\\n  --num_workers 4 \\\n  --size GPT2_large_div2_depth \\\n  --random \\\n  --resume final_runs/final_gpt2_large_div2_depth_check_bs512_lr0.0020_warmup3k_seqlen1024_debug/checkpoint-epoch=0-step=6499.ckpt \\\n  --doubling layers\n``` \n\n## Evaluation\n\nUse `evaluation/eval_wikitext.py` or `evaluation/eval_lambada.py` to evaluate [GPT-2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf) on one of the supported datasets. For example:\n\n```bash\npython evaluation/eval_wikitext.py\n```\n\nOr using Docker:\n\n```bash\ndocker build -t evaluation:latest .\ndocker run --rm --gpus all evaluation:latest evaluation/eval_wikitext.py\n```\n\n## Reference\n\nIf you use staged training in your research or wish to refer to the baseline results published here, \nplease use the following BibTeX entry. \n```\n@misc{shen2022staged,\n    title={Staged Training for Transformer Language Models},\n    author={Sheng Shen and Pete Walsh and Kurt Keutzer and Jesse Dodge and Matthew Peters and Iz Beltagy},\n    year={2022},\n    eprint={2203.06211},\n    archivePrefix={arXiv},\n    primaryClass={cs.CL}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fallenai%2Fstaged-training","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fallenai%2Fstaged-training","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fallenai%2Fstaged-training/lists"}