zimplex/producthelm-6x-checkpoints

ProductHELM 6× checkpoints

Full training checkpoints for ProductHELM-D with MiCE and TangentAttnRes, using two or four Lorentz factors. Both runs completed 4,372 optimizer updates and passed full-step, model-layout and finite model/optimizer-state validation.

Checkpoint	Lorentz factors	Architecture
producthelm_6x_f2_final.pt	2	L6_W2340_A36
producthelm_6x_f4_final.pt	4	L6_W2340_A36

Both models have six decoder blocks, total ambient width 2340, 36 attention heads and vocabulary size 128256. The first block uses a dense FFN; the other five use four routed experts, two activated experts and one shared expert. The configured expert intermediate width is 4680.

Training and data

The input was the unchanged prepared dataset achiverram28/producthelm-wikipedia-tokenized, at revision 8be3fbb638bc9b473d113173fad8f017c512ce6c: 6,715,364 rows. All 136 training Arrow files were verified, including token/mask bounds and checksums. The first 1,024 rows matched a canonical Wikipedia 20231101/Llama-3.1 preparation reference. The uploader supplied no additional preparation provenance card.

Each run used eight GPUs with ordinary data parallelism, one packed 2048-token sequence per GPU microstep and 16 accumulation steps: a global batch of 128 sequences. Training used BF16 autocast with FP32 parameters, seed 1234, AdamW plus RiemannianAdam, peak learning rate 4e-4, 3% warm-up, cosine decay to 10% of the peak, weight decay 0.01 and gradient clipping at 1.0.

Exact portable configurations are in training_configs/f2.json and training_configs/f4.json. File sizes, hashes and release provenance are in checkpoint_manifest.json. Evaluation metrics are not included in this release.

Load with the custom source

Use the original ProductHELM source and its dependencies. The source baseline is fd35e79. Eight-GPU training used the included 8_gpu_ddp.patch, which removes an expert-count/world-size divisibility assertion: ordinary DDP keeps all four routed experts replicated on every rank. The runtime used Python 3.10, PyTorch 2.7.1 with CUDA 12.8, Accelerate 1.3.0 and Geoopt 0.5.0.

Run this from the source checkout after installing its requirements:

import torch
from huggingface_hub import hf_hub_download
from runs import mice_args
from helm.hypercore.manifolds.product_lorentz import ProductLorentz
from helm.modules.product_helm_d import ProductLTransformerDecoder

factors = 2  # Set to 4 for the other checkpoint.
filename = f"producthelm_6x_f{factors}_final.pt"
path = hf_hub_download("zimplex/producthelm-6x-checkpoints", filename)
checkpoint = torch.load(
    path, map_location="cpu", weights_only=False, mmap=True
)
manifolds = [ProductLorentz(1.0, n_factors=factors) for _ in range(3)]
model = ProductLTransformerDecoder(
    *manifolds, arch="L6_W2340_A36", vocab_size=128256,
    context_length=2048, residual="tangentattnres", ffn="mice",
    moe_args=mice_args("6x"), n_dense_layers=1,
)
model.load_state_dict(checkpoint["model_state_dict"], strict=True)
model.eval()

Use the Llama-3.1 tokenizer NousResearch/Meta-Llama-3.1-8B. The custom decoder returns next-token logits through model(input_ids).

Checkpoint format

Each .pt file is the native full resume dictionary saved by save_checkpoint_both:

• global_step — 4372.
• model_state_dict.
• optimizer_state_dict — two optimizer state dictionaries.
• scheduler_euc_state_dict.
• scheduler_hyp_state_dict.

Configuration, dataset references and RNG state are not embedded in these files. Use the matching factor configuration and source when restoring them.

Licensing

The source code is MIT-licensed, copyright © 2025 Graph and Geometric Learning; the included patch retains that notice in LICENSE.source.txt. The source's vendored packing utility is Apache-2.0, as documented in its NOTICE. A separate license for these checkpoint weights has not been specified.