Picotron
About the project
Today we’re excited to share our open source project Picotron which won Best Hardware Adaptive Project (Hardware) @OpenAIBuildWeekDevpostHackathon 2023! We’ve created a “from scratch” LLM training framework that’s adaptive to your existing hardware resources.
Our aim with Picotron was to provide users with easy access to a toolchain where you could go from “I have an old/constrained GPU…” all the way up to “I can now train my own native LLM!” without needing to be overwhelmed by modern stack complexities of deep learning and its inherent difficulty in making changes, understanding what’s happening behind-the-scenes, or running it all on your personal accessible hardware.
The end result is a configurable native decoder model and training stack that will work well even if you’re running your model training on a Turing T4 with just fp16 but adapts accordingly if you’ve got more powerful hardware on hand.
What Picotron does
Picotron currently includes:
A native configurable decoder-only Transformer (PicotronDecoderModel) RoPE & per-layer NoPE positional schemes GQA + optional sliding-window attention Optional MoE & MLA model-paths CPU-safe config-validation & tests Single-GPU training / dist-train infrastructure Memory-mapped token-cache prep for HF-datasets Safetensors checkpoints + arch sidecar-configs picotron --config config.yaml CLI Rich/tqdm progress-bar with notebook-fallback CSV/Text run-logging Hardware aware precision (fp16 on T4 / bf16 if Ampere+) Optinal Triton RMSNorm/SwiGLU Kernels + PyTorch fallbacks Native checkpoint load + simple native-inference Experimental SFT-support for compat CLMs
Demo highlights:
Focuses on the most-important stable workflow right now: do we natively-pretrain on real fine-web-edu stuff? Can we use triton? Do checkpoints work? Does inference work?
And then I'm going to go ahead and try all these cool things! Let's see how they work together.
What inspired us
I’ve been working on this for a while, but I’ll describe how everything started here. As many of you know (or as many of you who are still reading would), there’s an alarming number of LLM training projects out there with two main flaws – they assume access to some sort of high end hardware, or else rely on a myriad of hidden implementations in their big dependency list, or then don’t surface errors at all if you mess up an optimization somewhere. Our previous prototype used Picotron had all these problems too 😅. dtype bugs in Triton, unsafe assumptions about CUDA behavior, checkpoint-resume issues… you name it!
This rebuild was therefore guided by one principle:
Correctness first of all. Each option should have its own clear, explicit and testable features (with a nice fallback). T4 shouldn't be considered an afterthought.
How we built it
It’s a fairly incremental approach to building Picotron – all of its pieces are modualar. Below we briefly describe each module:
• config/: Strict nested YAML validation • models/: Native decoder model • nn/: Attention, RoPE, MoE, MLA, & optional Triton kernels • data/: Synthetic data, packed token datasets, memmapped caches, & multi-dataset support • training/: Training loop, mixed precision, checkpoint integration, & logging • serialize/: Stores model weights as safetensors along with optimizer state & native architecture meta-data • parallel/: Distributed & ZeRO oriented infrastructural bits • picotron_sft/: Experimental scriptable layer for doing full fine-tuning
The native demo model above has a tied token/embedding/output table that keeps the number of parameters close to 1M but also works out the full language-model training pipeline. All of this is configured via YAML (and kicked off by the install cli).
So what exactly does this picotron command look like? If you’re on Kaggle, you’ll find:
picotron --config config_native_1m_fineweb.yaml
That will validate your installation of Triton (and whether you have the right cards) and then run our actual forward-and-backward probe on both RMSNorm and SwiGLU. At that point, assuming everything works, you’re free to start the longer running training job.
Challenges Faced
Now, I’m not going to pretend this wasn’t an epic struggle — but don’t think that means we were working on anything crazy. Writing the forward pass of the transformer is not actually one of them. What really took us forever was making sure all these things worked when subjected to true hardware constraints.
Triton Compatibility
It might seem simple, but just knowing that someone has the right Triton package isn’t enough. We had to make sure that once installed, they could compile (and execute!) their code successfully. This meant doing some deep diving into our probe kernel code during Kaggle testing… ultimately discovering two very important things:
We weren’t handling dtype APIs properly; There was some kind of SwiGLU lowering bug specific to Turing chips. Neither of which would’ve been discovered by simply trusting the optional nature of accelerations! So today Picotron explicitly probes and doesn’t even let you hit the long training phase until both RMSNorm/SwiGLU are good to go.
Training Costs
Even “small” sizes are subject to high costs, though! Just how small is our 1M parameter model? Well, remember all those tokens we need to project? Let’s calculate a few logit numbers:
logits_per_step = batch_size × sequence_length × vocabulary_size
With a GPT-2 sized vocabulary, there’s no way around it — even that tiny 1M model is going to spend a lot of time computing the big matrix of logits for each token. In short, relying on model size alone won’t cut it anymore — we needed to adjust the demo to provide realistic steps-per-epoch.
But what’s a demo without some cool looking logs? And yet Kaggle notebooks aren’t ordinary terminals. For maximum readability, rich displays show tables multiple times in capture mode, so Picotron checks for notebook contexts and falls back to safe tqdm outputs instead.
Checkpointing
You’d expect check pointing to be pretty basic stuff… But it turns out to be harder than expected. Not only did we save our native model parameters into safetensors format, we also stored the complete architecture configuration alongside the checkpoint so we knew all the required details to reconstruct the exact same model again (without having to memorize dimensions, etc.).
And lastly, it became abundantly clear during development that there were plenty of things we didn’t know ahead of time.
Lessons Learned
The most critical takeaway here? You have to pair hardware detection with proper runtime validation. “Optional” acceleration shouldn’t introduce unexpected correctness issues. Small model sizes ≠ low computational costs. Metadata matters as much as model parameters. And maybe the biggest lesson of all was realizing that building it ourselves forced us to consider certain performance, memory, and correctness trade-offs — many of which would be invisible otherwise!
What’s next for Picotron? Obviously, the best thing to do is keep developing it further — including additional real GPU verifications, trying out more attention backend options, improving native generation support with kv-cache features, and continuing to polish the various post-training experimental interfaces. As always, please don’t judge based off the video below — due to the constraints of hacking marathons, we decided it would be better to make a concise walk through of our successful native pretraining experience rather than showing you all possible errors we may have encountered.
(Also note that our Kaggle submission contains detailed documentation about the whole repo inside the README, as well as a public version of the GitHub Pages demo and lots more explanation of implementation choices, safety checks, and experimental features in our collection of Kaggle notebooks.)
Log in or sign up for Devpost to join the conversation.