aethertp

/

PicoLM-80M-Instruct

Like 2

Hi Andrew,

Thanks for checking out PicoLM-80M and sharing the Maba architecture!

Your analysis of the ARC-Easy floor and vocabulary parameter trade-off is spot
on. In our initial v1 run, we prioritized establishing a stable zero-budget
baseline on dual T4 GPUs, validating the QK-norm + GQA + WSD pipeline, and
verifying the tokenizer synchronization.

Interestingly, your first point aligns directly with our next iteration: we are
currently training PicoLM-V2, which implements immediate block-wise layer
sharing (MobileLLM-LS style) with 18 physical blocks executed twice to achieve
36 layers of computational depth, along with an expanded SwiGLU intermediate
dimension (1664) and targeted factual/science QA data injection to break the ARC
floor.

The low-rank vocabulary factorization idea 16384 to 64 to 576 is very
intriguing we will definitely evaluate it for our future parameter allocation
experiments.

I'll share the V2 empirical metrics as soon as the run concludes!

Best,
Emre

Hi Emre,

Fantastic to hear that! Converging on 18 physical blocks executed twice for 36 effective layers is a great architectural pivot for PicoLM-V2.

One quick practical observation if you are running MobileLLM-LS style weight sharing from scratch:
Watch the hidden state variance across the repeated passes. Because the second pass refines already-processed features rather than raw embeddings, representations tend to shift in magnitude. If you notice gradient instability during early warmup, keeping distinct RMSNorm gain parameters for each pass (Split RMSNorm) is usually the magic bullet that stabilizes the dynamics without adding perceptible parameter weight.

Expanding the SwiGLU width to 1664 alongside the targeted science QA mix should give ARC-Easy the exact relational boost it needs to break well past the 25% floor.

Really looking forward to seeing the V2 benchmarks once the run concludes. Best of luck with the T4 compute run!

Best,
Andrew

Hey Andrew,

The V2 run actually just wrapped up, and your intuition on the ARC floor was spot on.

We ran the evaluation on the base checkpoint using the same slices to keep the comparison fair:

• ARC-Easy: 25.60% (v1) -> 42.00% (+16.40% jump, finally out of the random floor)
• HellaSwag: 31.20% -> 34.40% (+3.20%)
• Validation Perplexity: 16.08 (on the new 24.5k vocab)

Expanding the SwiGLU width to 1664 along with the targeted QA data clearly made a massive difference for factual retrieval.

Also, thanks for the Split RMSNorm tip—that's a really clever way to handle representation shift in weight sharing without adding parameter weight. Fortunately, our QK-norm combined with Megatron scaling (scaled to 1/sqrt(2 * 36)) kept gradients smooth across the 11k steps, but I’m definitely keeping Split RMSNorm in mind for future iterations.

Currently running the SFT alignment for the instruct version. I'll drop an update here once the V2 model and weights are pushed to the hub!

Best,
Emre

Hi Emre,

A +16.4% jump on ARC-Easy to 42.00% is massive! Breaking cleanly out of the 25% random floor confirms that expanding effective computational depth via block sharing gives the network the multi-step relational capacity it was missing. HellaSwag picking up +3.2% alongside it confirms genuine generalization across domains.

Adjusting Megatron residual scaling to 1/sqrt(2 * 36) was a great engineering move. Dampening the residual branch mathematically to account for the 36-step computational path is the cleanest way to prevent activation explosion during backprop when reusing weights.

Speaking of next iterations, I just published the complete reference implementation and benchmarks for Maba v2 on Hugging Face (https://huggingface.co/AndrewThompson1233/maba-v2-architecture). Since you are training on Kaggle T4s in the 80M-100M envelope, it might be a really neat experimental setup to benchmark for PicoLM-V3:

1. Native T4 hardware optimization: The reference 101M config was profiled directly on Tesla T4, with fused Triton DGDA kernels hitting 264k-375k tok/s.

2. Vocabulary budget reclamation: Projecting your 24.5k vocab through a rank-128 bottleneck (24.5k -> 128 -> hidden_dim) drops the embedding tax down to ~4%, instantly reclaiming ~12M parameters that you can pour directly into deeper blocks or wider SwiGLU without leaving the 80M budget.

3. 3:1 Hybrid recurrence (DGDA + Latent MABA-SA): It keeps generation latency strictly flat at O(1) (35-37 ms/token on T4 even past 32k context) and slashes KV-cache by nearly 40x, which is huge for on-device deployment.

If you are curious to run a quick experimental comparison against pure attention on Kaggle to see how the hybrid recurrent dynamics affect convergence and ARC reasoning, the PyTorch code and Triton kernels are fully modular.

Really looking forward to the V2 weights dropping on the hub!

Best,
Andrew

Hey Andrew,

Thanks! The V2 instruct model, safetensors, and GGUFs are officially live on the hub now: https://huggingface.co/aethertp/PicoLM-V2-81M-Instruct

You were spot on about the autoregressive decoding dynamics with weight sharing. During our early SFT tests, the model was indeed slightly more prone to repetition loops if left unconstrained. Applying proper turn-by-turn prompt masking alongside a mild repetition penalty (around 1.15) stabilized the outputs and locked in clean, concise responses.

We also ran an interesting ablation: pushing a second epoch over the 379M dataset to reach ~757M tokens. Interestingly, while cross-entropy and validation perplexity kept improving monotonically (PPL dropped to 15.5), zero-shot downstream accuracy plateaued and showed a slight ~1-2% variance dip (ARC-Easy settled at 40.00%, HellaSwag at 33.20%). It was a great empirical confirmation that for sub-100M footprints, fresh token diversity strictly beats multi-epoch recycling.

Currently wrapping up a quick v2.1 alignment pass to sharpen algorithmic recursion and math scratchpads, and then planning out v3 where we'll definitely integrate your Split RMSNorm suggestion with an entirely fresh, non-repeating curriculum.

Would love to hear your feedback once you get a chance to take the V2 weights for a spin!

Best,
Emre

Hi Emre,

Took the PicoLM-V2 weights for a test run on Kaggle/CUDA as promised! The underlying checkpoint and the 18x2 layer-sharing dynamics look really promising, but running the model via standard AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True) hits several blockers in modeling_picolm_v2.py and tokenizer_config.json.

Here is a quick diagnostic and patch breakdown so you can push a fast hotfix to the repo:

1. Uninitialized RoPE buffers (Root cause of NaNs and CUDA asserts):
   In lines 92-93, self.register_buffer("rope_cos", cos, persistent=False) causes transformers to allocate uninitialized GPU memory (torch.empty) when instantiating from the Hub. Because the buffers are neither in the checkpoint nor recomputed in forward(), rope_cos contains garbage floats (e.g. -4.59e36 / NaNs), causing immediate device-side asserts during sampling.
   Fix: Either recompute RoPE freqs dynamically inside forward() based on sequence length, or ensure the buffer is populated in __init__ and properly registered.

2. Attention dtype mismatch in SDPA:
   RoPE freqs are calculated in float32. Multiplying q and k by cos/sin automatically upcasts them to float32, while v remains float16. F.scaled_dot_product_attention then throws a fatal dtype mismatch error.
   Fix: Ensure cos and sin are cast to q.dtype before rotary application (cos.to(q.dtype)).

3. Attribute typo in out_proj:
   In line 41, self.out_proj = nn.Linear(args.n_heads * self.head_dim, ...) throws an AttributeError because your config uses num_attention_heads (or config.n_heads).

4. Missing GenerationMixin for .generate():
   In recent transformers releases (4.45+), PreTrainedModel no longer inherits directly from GenerationMixin. If the model class only inherits from PreTrainedModel, calling .generate() fails with an AttributeError.
   Fix: Define the class as class PicoLMV2ForCausalLM(PreTrainedModel, GenerationMixin):.

5. Forward signature and KV-cache handling:

• Add attention_mask: Optional[torch.Tensor] = None to the forward() signature to prevent transformers kwargs validation errors.
• Ensure all_tied_weights_keys = ["lm_head.weight"] is declared so transformers handles the tied embedding head without warnings.
• Currently, generating with use_cache=True slices input_ids to the trailing token, but without internal past_key_values logic in the attention blocks, the model loses conversational history. Setting use_cache=False as default in generation config is recommended until the cache mechanism is wired up.

6. Missing chat_template in tokenizer_config.json:
   tokenizer.apply_chat_template() throws a ValueError because the Jinja template is missing from tokenizer_config.json. Adding a standard ChatML or Llama-style template string will make the snippet in your model card work out of the box.

By the way, absolutely no pressure at all and totally fine if you'd rather not, but since our discussions and tips on block recycling, computational depth scaling (MobileLLM-LS style), and residual dampening helped shape the V2 jump (+16.4% on ARC-Easy), would you be open to adding a brief shoutout or Acknowledgments link to Maba v2 (https://huggingface.co/AndrewThompson1233/maba-v2-architecture) in the PicoLM-V2 model card? Completely up to you either way - just thought I'd ask since it would really help an indie architecture project get discovered! :)

Best of luck with the v2.1 patches and v3 planning!

Best,
Andrew