OLMo-core + Engram graft: small-scale debug comparison

A debug comparison between a base OLMo3 600M model and an Engram memory variant showed the grafted model achieved lower training and evaluation cross-entropy loss and faster gradient norm stabilization, indicating successful integration and improved early learning behavior.

I ran a 200-step , with global batch size =32 debug comparison between a base OLMo3 600M model and the same dense backbone with a DeepSeek-style Engram memory graft. The goal was to check whether the custom module was wired correctly, whether FSDP/HSDP wrapping and optimizer handling were stable, and whether the training/eval curves looked coherent. Base model: Engram variant: ~1.7B trainable parameters Engram injected into layers 1 and 5 Most added parameters come from sparse/hash-memory capacity, so total parameter count is not an apples-to-apples proxy for dense active compute. Both are trained with Dion https://github.com/microsoft/dion/ https://github.com/microsoft/dion/ optimizer. Under the same short debug setup, the Engram variant showed: lower train CE loss lower eval CE loss / PPL slightly faster grad-norm stabilization The early signal is encouraging: the Engram graft is training-shaped, stable, and appears to improve early learning behavior in this setup. Custom architecture work is not just “does the forward pass run?” For this integration, the parameter hierarchy, wrapping policy, optimizer handling, memory profile, and training curves all had to line up. Earlier versions trained mathematically, but had poor memory behavior because the custom modules were not placed inside the wrapped block hierarchy. W&B logs: Weights & Biases https://wandb.ai/jenwei0312/olmo3-engram-experiments