{"id":19072286,"url":"https://github.com/andron00e/learning-at-scale","last_synced_at":"2026-03-07T08:33:23.591Z","repository":{"id":259045560,"uuid":"876186967","full_name":"Andron00e/learning-at-scale","owner":"Andron00e","description":"A codebase for training models of different scales","archived":false,"fork":false,"pushed_at":"2025-12-30T17:53:22.000Z","size":461,"stargazers_count":5,"open_issues_count":0,"forks_count":2,"subscribers_count":2,"default_branch":"main","last_synced_at":"2026-01-03T08:56:45.882Z","etag":null,"topics":["llms","pretraining"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/Andron00e.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2024-10-21T14:41:31.000Z","updated_at":"2025-12-30T17:53:26.000Z","dependencies_parsed_at":"2024-12-03T01:27:42.317Z","dependency_job_id":"540f4315-691d-417e-82f6-293615e29660","html_url":"https://github.com/Andron00e/learning-at-scale","commit_stats":null,"previous_names":["andron00e/learning-at-scale"],"tags_count":0,"template":true,"template_full_name":null,"purl":"pkg:github/Andron00e/learning-at-scale","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Andron00e%2Flearning-at-scale","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Andron00e%2Flearning-at-scale/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Andron00e%2Flearning-at-scale/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Andron00e%2Flearning-at-scale/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/Andron00e","download_url":"https://codeload.github.com/Andron00e/learning-at-scale/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/Andron00e%2Flearning-at-scale/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":30209942,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-03-07T05:23:27.321Z","status":"ssl_error","status_checked_at":"2026-03-07T05:00:17.256Z","response_time":53,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"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":["llms","pretraining"],"created_at":"2024-11-09T01:36:30.424Z","updated_at":"2026-03-07T08:33:23.571Z","avatar_url":"https://github.com/Andron00e.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Learning @ Scale: a modular codebase to experiment with different ML model training settings\n\nThis repository is designed to help researchers create reproducible experiments. Inspired by the [epfml/llm-baselines](https://github.com/epfml/llm-baselines) codebase and the [nanoGPT](https://github.com/karpathy/nanoGPT) benchmark.\n\n## Quickstart \n\nInstall dependencies: \n\n```\npip install -r requirements.txt\n```\n\nRun a simple training on the Slimpajama dataset ([6B subset](https://huggingface.co/datasets/DKYoon/SlimPajama-6B), 24GBs decompressed, takes a few minutes to download):\n\n```sh\npython ./src/main.py --config_format base\n```\n\nThe above command trains a 123.59M parameters model. It trains for 25k iterations with a batch size of 128=32x4 (4 gradient accumulation steps), using a cosine schedule with a maximum learning rate of 1e-3 that is reduced to 1e-4 at the end of training. The model is saved in the `./exps` folder.\n\nThis training takes roughly ~3h on a single A100 (80GB) GPU. The plot of the training and validation loss should look roughly like this:\n\n\u003cimg src=\"./assets/loss_slimpajama.png\" alt=\"Loss on SlimPajama\" width=\"500\"/\u003e\n\u003cimg src=\"./assets/pplx_slimpajama.png\" alt=\"Perplexity on SlimPajama\" width=\"500\"/\u003e\n\nYou can check out the wandb run for yourself [here](https://wandb.ai/haeggee/llm-lauzhack/runs/lm2obqy9?nw=nwuserhaeggee).\n\n\n## Less quick start\n\nHere are the possible parameters you can use (copypasted from `config/base.py`):\n\n```python\n# General training params\nparser.add_argument('--batch_size', default=32, type=int)\nparser.add_argument('--acc_steps', default=4, type=int)\nparser.add_argument('--seed', default=0, type=int) # random seed for the parameters\nparser.add_argument('--data_seed', default=1337, type=int) # random seed defining the data ordering\nparser.add_argument('--eval_interval', default=200, type=int)\nparser.add_argument('--full_eval_at', nargs=\"+\", type=int)\nparser.add_argument('--eval_batches', default=32, type=int)\nparser.add_argument('--device', default='cuda:0', type=str) # see below to run on multiple GPUs\nparser.add_argument('--iterations', default=25000, type=int) # total number of training iterations\nparser.add_argument('--warmup_steps', default=300, type=int)\nparser.add_argument('--lr', default=1e-3, type=float)\nparser.add_argument('--wsd_final_lr_scale', default=0.0, type=float) # wsd scheduler\nparser.add_argument('--wsd_fract_decay', default=0.1, type=float) # wsd scheduler \nparser.add_argument('--decay_type', default='linear', choices=['linear', 'cosine', 'exp', 'miror_cosine', 'square', 'sqrt'])\nparser.add_argument('--dd_second_decay_type', default='linear', choices=['linear', 'cosine', 'exp', 'miror_cosine', 'square', 'sqrt'])\nparser.add_argument('--dd_first_lr_factor', default=1e-2, type=float)\nparser.add_argument('--weight_decay', default=0.1, type=float) # I recommend you keep this value, else instabilities might arise\nparser.add_argument('--beta1', default=0.9, type=float) # adam parameter\nparser.add_argument('--beta2', default=0.95, type=float) # adam parameter\nparser.add_argument('--scheduler', default='cos', choices=['linear', 'cos', 'wsd', 'cos_inf', 'none', 'dd'])\nparser.add_argument('--cos_inf_steps', default=0, type=int) # cos_inf scheduler\nparser.add_argument('--opt', default='adamw', choices=['adamw', 'sgd', 'muon', 'soap', 'ademamix', 'ademamix2', 'lion', 'sf-adamw', 'sf-sgd', 'signsgd', 'signum', 'sgdf', 'prodigy', 'sophiag', 'shampoo', 'adopt', 'clip-adagrad', 'clip-adagrad-delay-eta', 'clip-adam', 'clip-adam-delay-eta', 'mars', 'adafactor', 'lamb'])\nparser.add_argument('--eval_freq', default=200, type=int) # in iterations\nparser.add_argument('--results_base_folder', default=\"./exps\", type=str) # where the checkpoints will be saved\nparser.add_argument('--grad_clip', default=0.0, type=float) # default value is 1.0 in NanoGPT\nparser.add_argument('--momentum', default=0.9, type=float)\nparser.add_argument('--shampoo_beta', default=-1.0, type=float)\nparser.add_argument('--precondition_frequency', default=10, type=int)\nparser.add_argument('--max_precond_dim', default=10000, type=int)\nparser.add_argument('--merge_dims', default=False, type=bool) # merge dimensions till the product of the dimensions is less than or equal to max_precond_dim\nparser.add_argument('--precondition_1d', default=False, type=bool)\nparser.add_argument('--normalize_grads', default=False, type=bool)\nparser.add_argument('--soap_data_format', default='channels_first', type=str)\nparser.add_argument('--correct_bias', default=True, type=bool)\nparser.add_argument('--nesterov', default=False, type=bool) # whether to use Nesterov-style momentum \nparser.add_argument('--muon_ns_steps', default=5, type=int) # the number of steps to use in the newton schulz, if it is iterative\nparser.add_argument('--muon_lr_factor', default=0.02, type=float) # a factor by which to reduce the lr for muon\nparser.add_argmunet('--adema_beta3', default=0.9, type=float) # beta3 in AdEMAMix\nparser.add_argument('--adema_alpha', default=2.0, type=float) # alpha in AdEMAMix\nparser.add_argument('--adema_beta3_warmup', default=None, type=int) # AdEMAMix hyperparameter\nparser.add_argument('--adema_alpha_warmup', default=None, type=int) # AdEMAMix hyperparameter\nparser.add_argument('--schedulefree_r', defalut=0.0, type=float) # schedulefree hyperparameter\nparser.add_argument('--weight_lr_power', default=2.0, type=float) # schedulefree hyperparameter\nparser.add_argument('--model_sharding', default=None, type=bool) # Adam-mini\nparser.add_argument('--adam_mini_verbose', default=False, type=bool) # print all the logs if true\nparser.add_argument('--log_interval', default=50, type=int)\nparser.add_argument('--dampening', default=0.0, type=float)\nparser.add_argument('--prodigy_beta3', default=None, type=float) # coefficients for computing the Prodidy stepsize using running averages\nparser.add_argument('--prodigy_decouple', default=True, type=bool) # Use AdamW style decoupled weight decay\nparser.add_argument('--prodigy_use_bias_correction', default=False, type=bool)\nparser.add_argument('--prodigy_safeguard_warmup', default=False, type=bool) # Remove lr from the denominator of D estimate to avoid issues during warm-up stage. Off by default.\nparser.add_argument('--prodigy_fsdp_in_use', default=False, type=bool)\nparser.add_argument('--sophia_rho', default=0.04, type=float)\nparser.add_argument('--clipping_type', default='no', choices=['no', 'local', 'elementwise']) # for methods with clipping\nparser.add_argument('--clipping_eta', default=1.0, type=float)\nparser.add_argument('--mars_type', default='mars-adamw', choices=['mars-adamw', 'mars-lion', 'mars-shampoo'],)\nparser.add_argument('--mars_vr_gamma', default=0.025, type=float)\nparser.add_argument('--mars_is_approx', default=True, type=float)\nparser.add_argument('--mars_lr', default=3e-3, type=float)\nparser.add_argument('--mars_beta1', default=0.95, type=float)\nparser.add_argument('--mars_beta2', default=0.99, type=float)\nparser.add_argument('--adafactor_decay_rate', default=-0.8, type=float)\nparser.add_argument('--lamb_use_bias_correction', default=False, type=bool)\n# Dataset params\nparser.add_argument('--dataset', default='slimpajama', choices=['slimpajama', 'wikitext', 'shakespeare-char', 'arxiv', 'arxiv2000', 'arxiv+wiki', 'openwebtext2', 'redpajama', 'redpajamav2', 'slimpajama_chunk1', 'fineweb', 'finewebedu'])\nparser.add_argument('--tokenizer', default='gpt2', type=str, choices=['gpt2', 'mistral'])\nparser.add_argument('--vocab_size', default=50304, type=int)\nparser.add_argument('--data_in_ram', action='store_true') # force the data to RAM, you most likely do not need this  \n# Model params\nparser.add_argument('--model', default='base', choices=['base', 'llama', 'test'])\nparser.add_argument('--parallel_block', action='store_true')\nparser.add_argument('--use_pretrained', default='none', type=str) # 'none', 'gpt2' or a path to the pretraind model\nparser.add_argument('--from_dense', action='store_true')\nparser.add_argument('--init_std', default=0.02, type=float)\nparser.add_argument('--dropout', default=0.0, type=float) # keep to 0 unless in low data regime (e.g. wikitext)\nparser.add_argument('--n_head', default=12, type=int)\nparser.add_argument('--n_layer', default=12, type=int) # depth in (att + ff) blocks\nparser.add_argument('--n_embd', default=768, type=int) # hidden size ... \nparser.add_argument('--sequence_length', default=512, type=int)\nparser.add_argument('--dtype', default='bfloat16', type=str, choices=['float32', 'float16', 'bfloat16'],)\nparser.add_argument('--bias', default=False, type=bool)\nparser.add_argument('--compile', action='store_true') # if true then model is compiled \nparser.add_argument('--rmsnorm_eps', default=1e-5, type=float) # used by the llama model\nparser.add_argument('--multiple_of', default=256, type=int) # used by the llama model make SwiGLU hidden layer size multiple of large power of 2\nparser.add_argument('--n_kv_head', default=None, type=int) # for Adam-mini\n# Checkpointing\nparser.add_argument('--results_base_folder', default='./exps', type=str)\nparser.add_argument('--permanent_ckpt_interval', default=0, type=int)\nparser.add_argument('--latest_ckpt_interval', default=0, type=int)\nparser.add_argument('--resume_from', default=None, type=str)\nparser.add_argument('--resume_from_swa', default=None, type=str)\nparser.add_argument('--auto_resume', default=True)\n# logging params (WandB)\nparser.add_argument('--wandb', action='store_true') # whether to use wandb or not\nparser.add_argument('--wandb_project', default='my-project', type=str)\nparser.add_argument('--wandb_entity', default=None, type=none_or_str) # for the team projects\nparser.add_argument('--wandb_run_prefix', default='none', type=str) # is added before the autogenerated experiment name\nparser.add_argument('--eval_seq_prefix', default=\"Once upon a time\", type=str) # prefix used to generate sequences\nparser.add_argument('--log_dynamics', action='store_true')\n# Distributed args\nparser.add_argument('--distributed_backend', default=None, type=str, required=False,\n                    choices=distributed.registered_backends())  # distributed backend type (e.g. nccl)\n```\n\n## Using WandB\n\nYou need to give your wandb authorize key in order to send the data to your wandb account. If you start jobs on a server without access to prompt, then you can set the `WANDB_API_KEY` variable within your script:\n\n```bash\n# this is a script that could be executed on a server\npip install -r requirements.txt # install req.\nexport WANDB_API_KEY=\"put your authorize key here, to find it: https://wandb.ai/authorize\"\npython ./src/main.py --config_format base --wandb --wandb_project \"my awesome project\" --n_layer 7 --model base --seed 123\n```\n\n## Want to contribute? \n\nFeel free to add a new architecture, algorithm, dataset, or just modify our code.\nThe structure of the project is the following: \n\n```sh\nsrc/\n    main.py         # pick the right data, model, and training function\n    config/\n        __init__.py # contains CONFIG_FORMAT_TO_MODULE_MAP mapping the name given to the --config_format flag with a python conf file\n        base.py     # config for the base model\n    data/\n        utils.py    # contains the get_dataset function\n        wikitext.py # load/process wikitext\n        arxiv.py    # load/process arxiv\n        fineweb.py # load/process the Fineweb dataset\n        slimpajama.py\n        ...\n    models/\n        utils.py    # contains the get_model function\n        base.py     # contains the standard transformer base architecture\n        llama.py    # llama architecture\n    optim/\n        utils.py    # contains eval and get_batch functions\n        base.py     # training function for the base and llama models\n        ...\n    distributed/\n        # code to enable simple distributed training\n```\n\n* Given the above structure, to add your own model, you can just fork this repository and modify the `./src/models/base.py` file, then if necessary update `./src/optim/base.py` in case you need some custom training loop or evaluation. You also need to fork the `./src/config/base.py` file to add your own parameters, which imply adding your new config to the mapping `CONFIG_FORMAT_TO_MODULE_MAP` in `./src/config/__init__.py`. \n* To add a new dataset, create a new file in the `./src/data` folder, take a look at `wikitext.py` for the expected format. \n* For optimizers, store all code inside the `./src/optim/` directory; it would be convenient if you simply add one file `yet_another_method.py`. Make sure to reference you optimizer with its adjusted arguments in `./src/main.py`, `./src/config/base.py` and `README.md` files of your branch/fork.\n\n## Codestyle\n\nWe  use [black](https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html) and [isort](https://pycqa.github.io/isort/) for all pull requests. Before committing your code, simply run ```black . \u0026\u0026 isort .```\n\n## Multi-GPU training\n\nGiven a multi-GPU machine with e.g. 4 GPUs, one can distribute the training using data-parallelism:\n\n```sh\ntorchrun --nproc_per_node=4 ./src/main.py --config_format base --distributed_backend nccl --dataset slimpajama --model base\n```\n\nWhen using multiple GPUs, the data will be distributed among the GPUs by dividing the number of accumulation steps by the number of nodes. For instance if we train with a batch size of 32 and 4 accumulation steps, then each GPU will process batches of 32 elements and do 1 accumulation steps. For this reason we require `acc_steps` to be a multiple of the number of GPUs.    \n\n\n## Experimenting locally on your device with CPU\nIf do not have access to a GPU or just want to try the code locally on your device, you can try the Shakespeare dataset with character-level tokens:\n\n```sh\npython ./src/main.py --n_layer=2 --n_head=4 --n_embd=128 --sequence_length=256 --dataset=shakespeare-char --device=cpu --vocab_size=96\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fandron00e%2Flearning-at-scale","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fandron00e%2Flearning-at-scale","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fandron00e%2Flearning-at-scale/lists"}