a2d-gpt-neox-400M

<!-- This model card has been generated automatically according to the information the Trainer had access to. You should probably proofread and complete it, then remove this comment. -->

a2d-gpt-neox-400M

Base autoregressive → MLM diffusion models derived from the EleutherAI Pythia family.
These are base models only (no supervised fine-tuning / instruction tuning).

Model description

• Model type: Decoder-only Transformer converted to Masked Language Modeling (MLM) diffusion architecture
• Base architecture: Pythia (GPT-NeoX)

Variants

Original Pythia pretraining & checkpoint selection (EleutherAI):
All models were initialized from Pythia checkpoints selected according to Chinchilla scaling laws (Hoffmann et al., 2022) targeting ~15 tokens per parameter (slightly undertrained regime) for consistency across the model family.

• Dataset: The Pile
• Total tokens in Pythia run: 299.9B
• Checkpoint interval: every 2,097,152 tokens (~2M batch size)

Selected checkpoints:

• 70M: step 512 (~1.07B tokens, ~15 tok/param)
• 160M: step 1000 (~2.1B tokens, ~13 tok/param)
• 410M: step 3000 (~6.29B tokens, ~15.5 tok/param)

Intended uses & limitations

• Research on autoregressive-to-diffusion conversion
• Continued pre-training and domain adaptation experiments
• Analysis of scaling behavior in hybrid MLM diffusion models
• Not intended for direct deployment or chat/instruction use (base models only)

Training and evaluation data

Continued pre-training: All models were further trained on the same subset of ~3B tokens sampled from FineWeb-Edu using (masked) causal language modeling objective.

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 0.0001
• train_batch_size: 48
• eval_batch_size: 48
• seed: 46
• optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.85,0.94) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
• lr_scheduler_type: cosine
• lr_scheduler_warmup_steps: 625
• training_steps: 6250

How to Get Started with the Model

Use the following code to get started with the model:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "avgJo3/a2d-gpt-neox-160M"  # or 70m / 160m variant
revision = "main"  # or specific checkpoint tag if uploaded separately

model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    revision=revision, 
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
    model_name, 
    revision=revision, 
    trust_remote_code=True
)

Framework versions

• Transformers 5.0.0
• Pytorch 2.10.0+cu128
• Datasets 4.8.5
• Tokenizers 0.22.2

Citation

@article{biderman2023pythia,
  title={Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling},
  author={Biderman, Stella and others},
  journal={ICML},
  year={2023}
}