DeepSeek Engram × OLMo-core: Distributed Implementation

A documentation for my experimental integration of DeepSeek’s Conditional Memory via Scalable Lookup (Engram) architecture into AI2’s OLMo-core distributed training framework.

Status: Proof of Concept (Compiles & Trains)

Compute: 4x NVIDIA A40 (RunPod)

Framework: PyTorch 2.6.0 / OLMo-core

🧠 The Architecture

This project replaces standard Feed-Forward Networks (FFNs) or traditional Mixture-of-Experts (MoE) with an Engram memory layer.

By mapping the concepts from the DeepSeek paper to OLMo’s codebase, this architecture leverages:

• Token Routing: Dynamically routing tokens to specific Engram blocks rather than activating the entire memory matrix.
• 1D / 2D Block Parallelism: Designed to scale across multiple GPUs, minimizing communication overhead during the lookup phase.
• Compatibility: Native integration with OLMo-core’s HSDP (Fully Sharded Data Parallel) and torchrun distributed topologies.

🛠️ The RunPod MLOps Survival Guide

AI2’s OLMo-core is a beautiful, highly optimized framework, but it is deeply coupled with their internal supercomputer (Beaker / Gantry). Running this on independent cloud compute (like RunPod) requires surgically bypassing corporate cluster management and avoiding massive I/O bottlenecks.

1. The 90-Second “Burn & Rebuild” Pipeline

Do not use RunPod Persistent Network Volumes to install the environment. The I/O latency when compiling C++ / CUDA kernels like flash-attn will stall the pod for hours.

Instead, allocate a large NVMe Container Disk (e.g., 100GB) and build the environment ephemerally using the provided set_env.sh script.

# Clone the repo and run the setup
git clone b feature/engram-poc https://github.com/JenWei0312/OLMo-core.git
cd OLMo-core
source set_env.sh

2. Bypassing AI2 Beaker/Gantry

To train outside of AI2’s cluster, we must strip out their proprietary monitoring telemetry. The set_env.sh script automatically applies a sed sledgehammer to bypass the GantryCallback decorator:

sed -i 's/.*@GantryCallback.register.*/# Bypassed Beaker/g' src/olmo_core/launch/beaker.py

3. The Dataset Trap

AI2’s default dataloaders point to 200GB+ training chunks. If you are running a single-node smoke test, the script will silently hang while downloading. We generate raw binary dummy data locally (.tofile(), bypassing numpy ASCII headers which crash the C++ dataloader) to verify gradient flow before paying for S3 egress.

📊 Proof of Life

Distributed training successfully validated on a 4x A40 cluster using a 3.5B parameter (OLMo 3B + Engram) configuration.

• Throughput: ~36 Tokens / Second / Device (Steady state)
• Memory: 28.6 GiB Reserved / 24.0 GiB Active (Highly efficient utilization)
• Gradients: Stable. train/CE loss dropped from 12.21 → 12.05 over 196 steps with zero NaN spikes or deadlocks.

🚀 Next Steps (The Ablation Study)

With the infrastructure pipeline solidified and gradient stability proven on dummy data, the next phase is a full ablation study across a 2x2 architecture grid using real Dolma datasets. (Compute sponsorship pending).

🔗 Link to Repo

https://github.com/JenWei0312/OLMo-core/tree/feature/engram-poc