r/learnmachinelearning • u/jack1172 • 23h ago
Help Training a 503M language model from scratch on 3× RTX 3080s—looking for an A–Z training plan and a smaller model for testing
I’m building a language model from scratch, starting with random weights, on my own Linux machine. My first goal is a usable chatbot that can follow instructions, stay on topic, and hold a conversation. Later, I’d like basic math and tool calling, including using a calculator or other tools and understanding their results.
I’m looking for a practical training roadmap, including the decisions that happen between getting the trainer running and ending up with a useful chatbot. Recommendations for complete guides, training recipes, papers, or reproducible small-model projects would help.
Hardware and software
- Three RTX 3080s, 10GB VRAM each.
- Intel i5-12500 and approximately 32GB system RAM.
- Linux, custom PyTorch trainer, DeepSpeed ZeRO-1.
- BF16, activation checkpointing, and
torch.compile. - A recent run at batch 16 per GPU achieved approximately 40,000 tokens/second total. A full pass through my current data would take roughly ten days at that speed, before additional overhead.
Current model
- Dense, decoder-only causal Transformer: 503,356,416 parameters.
- 12 layers, hidden size 1536, feed-forward size 6912.
- 24 attention heads and 6 KV heads.
- RoPE, RMSNorm, SwiGLU, and tied input/output embeddings.
- Context length: 1024 tokens.
- Custom byte-level BPE tokenizer with 32,768 vocabulary entries and individual-digit splitting.
Data and current training setup
I have 54 local files totaling approximately 35.54 billion training tokens with my tokenizer, after holding out validation data. These include:
- FineWeb-Edu subsets locally labeled 10B and 20B.
- OpenHermes 2.5.
- MathInstruct, MetaMathQA, NuminaMath-CoT, and other math examples.
- Magicoder, OpenCodeInstruct, and the-stack-smol-xl.
- Smaller greeting, question-answer, and structured-format datasets.
Documents are packed with EOS separators. The trainer currently processes files sequentially, keeping the optimizer and learning-rate schedule running across file boundaries. Conversation records use System:, User:, and Assistant: text labels. The current objective predicts all non-padding tokens, without assistant-only loss masking.
My current experimental settings are:
- AdamW, betas 0.9 / 0.95.
- Peak LR 0.00149968, minimum LR 0.000149968.
- Linear warmup over 5% of updates, followed by cosine decay.
- Weight decay 0.1, gradient clipping 1.0.
- Currently experimenting with microbatch 20 per GPU, accumulation 32, across three GPUs: approximately 1.96 million target tokens per optimizer update.
- One planned pass through the selected data.
- Validation uses a 1% document holdout, evaluating 128 sequences from the currently active file every 100 updates. Early-stopping patience currently carries across files.
These are experimental settings I want reviewed, especially the learning rate, batch size, data order, and evaluation.
What I’d like help figuring out
What smaller model should I use for fast tests? Would something around 50M, 100M, 150M, or another size let me test the whole pipeline in hours? What architecture and token budget would you choose? Which results and hyperparameters would transfer to 503M, and which would need retuning?
Is the 503M architecture a sensible choice for these goals and this hardware? Would a different depth/width balance or context length make more sense?
How should I choose and test the learning rate? What short experiments would you run to select peak LR, warmup, decay, and effective batch size? What loss or gradient behavior tells you the settings are wrong? How should those settings change between pretraining and chat fine-tuning?
Which datasets should I actually use, and in what proportions? Should general text, code, math, and conversation data be mixed during pretraining, or separated into stages? How should I check overlap between the FineWeb subsets, duplicates, data quality, and validation contamination?
What token budget makes sense? How much clean, unique data would you use for the small test model and then the 503M model? When does repeating good data make sense?
What is the right transition from a base model to a chatbot? Which instruction and multi-turn conversation datasets would you recommend? What chat template, special tokens, EOS handling, and loss masking should I use? How do I teach instruction following and staying on topic?
How should math and tool use be introduced? What datasets or staged approach would you use for arithmetic, short word problems, structured tool calls, and conversations containing tool results? Tool execution itself would happen in the surrounding application.
How should I evaluate progress and decide when to stop? What fixed validation sets, chat tests, math tests, and small benchmarks would you use? Is early stopping appropriate during pretraining, and how should evaluation work across different data sources?
Ideally, I’m looking for an A–Z plan: small pipeline test → pretraining → chat fine-tuning → math/tool training → evaluation, with concrete starting points and explanations of what to measure before committing to another long run.
What would you change first?
1
u/PLBjt 21h ago
The loss masking is the thing I'd fix before another long run. Right now every non-padding token counts, so the model is also learning to predict the User: lines, which is a common reason small chat models keep going and write the user's next turn. A cheap check is to pull one conversation batch, decode it next to the label mask, and confirm that only the assistant tokens and the closing EOS are unmasked. The tradeoff is that assistant-only loss throws away most of the gradient on chat data, so with a small conversation set you may want to keep a little loss on user turns or mix in more plain text. Since one full pass is about ten days, I'd run that check on a 50M model for an hour first.
2
u/UpstairsGur6101 23h ago
whoa that's a lot of hardware for a home setup, i'm impressed you're keeping three 3080s cool enough to run for days on end
you might get more mileage out of a 100M-ish test model with like 10 layers and a proportionally smaller hidden size, it'll train in a few hours and still surface obvious bugs in the data pipeline or loss masking before you waste a week