The original Nano effective batch was 12 micro-batches x 11 accumulation
steps = 132 sequences. This distributed run preserves that exact global batch.
Ranks receive 16 or 17 sequences and scale their local mean losses so DDP's
averaged gradient is the true 132-sequence global mean.
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model_id = "Banaxi-Tech/muon-model-test"
5tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained(
7 model_id,
8 trust_remote_code=True,
9 dtype=torch.bfloat16,
10 device_map="auto",
11)
This is a base model, not an instruction-tuned chat model.