train-llm-from-scratch

Strip both DDP and torch.compile wrapper prefixes from every state-dict key, repeatedly, because they can be nested in any order. Replace the single module. check with:

def _strip_ddp_prefix(state_dict: dict) -> dict:
    out = {}
    for k, v in state_dict.items():
        while k.startswith("module.") or k.startswith("_orig_mod."):
            k = k.removeprefix("module.").removeprefix("_orig_mod.")
        out[k] = v
    return out

Use this same function in all checkpoint-loading paths, including inference and evaluation:

from src.post_training.utils import _strip_ddp_prefix

state = ck["model_state_dict"] if "model_state_dict" in ck else ck
state = _strip_ddp_prefix(state)
model.load_state_dict(
    {k: v for k, v in state.items() if k in model.state_dict()},
    strict=False,
)

For new checkpoints, save the unwrapped model so keys are clean from the start. The loader should also treat missing expected parameters as an error to prevent accidentally running with randomly initialized weights. After the fix, compiled/DDP checkpoints should load with missing: 0 keys.

Use the repository's CPU-friendly laptop/tiny preset. First prepare the TinyStories dataset and train a BPE tokenizer, then train the 369K-parameter model, and finally generate text:

python scripts/prepare_tiny_data.py
python scripts/train_transformer.py --preset tiny --arch modern
python scripts/generate_text.py --model_path models/tiny.pt

The tiny preset is designed for CPU-only machines and takes about 5 minutes on a typical laptop CPU. Larger student and small presets are also available. Run python scripts/benchmark.py to check training time on your hardware, and see docs/student for the guided learning path.

The OOM is expected because the default config uses roughly 3B parameters (N_EMBED=2048, N_BLOCKS=64), and optimizer state plus activations exceed 40GB. Update to the version containing PR #29 and run training with the opt-in memory optimizations:

python scripts/train_transformer.py --amp --amp-dtype bf16 --grad-checkpointing --grad-accum 8

To estimate VRAM usage before starting training:

python scripts/train_transformer.py --report-memory

--amp enables mixed precision, --grad-checkpointing reduces activation memory by recomputing transformer blocks, and --grad-accum 8 simulates a larger batch using smaller micro-batches. These changes allow the 3B config to fit on an A100 40GB and can also reduce VRAM usage on RTX 3090/4090/5090-class GPUs. If fragmentation is reported, also try:

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python scripts/train_transformer.py --amp --amp-dtype bf16 --grad-checkpointing --grad-accum 8

Checkpoint resume support is available on main in scripts/train_transformer.py. During training, enable periodic checkpoints with --checkpoint-every N; use --keep-last K to retain only the most recent checkpoints. After an interruption, restart with --resume latest to load the newest checkpoint, or --resume /path/to/checkpoint.pt for a specific one. The restored checkpoint includes model weights, optimizer state, LR scheduler state, loss history, and the last completed step, so training continues from the next step. Example: python scripts/train_transformer.py --checkpoint-every 1000 --keep-last 3 --resume latest. Checkpoints are written atomically using a temporary file plus os.replace to reduce corruption risk if the process crashes while saving.

The 404 means the Google Drive file was deleted or its sharing permissions are private. Remove the broken link from README.md, or replace it with a valid shared link after enabling link sharing in Google Drive. In the reported case, the fix was to remove https://drive.google.com/file/d/1Q_iklJ1RVGSb-Pdey8BHy3k8IF3UJv0z/view and keep the GitHub reference instead. Commit the change, for example: git add README.md && git commit -m "fix: remove broken resume Google Drive link".