Upload step_002000 checkpoint, training script and run command.

RUN_COMMAND.txt

@@ -0,0 +1 @@
+nohup python /workspace/nanomind/train.py --data_path /workspace/nanomind_data/pretrain_1m.jsonl.gz --out_dir /workspace/nanomind_runs/run1 --tokenizer_name hf-internal-testing/llama-tokenizer --seq_len 2048 --hidden_size 512 --n_layers 16 --n_heads 8 --n_kv_heads 1 --global_batch_size 64 --micro_batch_size 1 --lr 1e-3 --warmup_steps 2000 --max_steps 50000 --save_every 1000 --bf16 > /workspace/nanomind_runs/run1/train.log 2>&1

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 1126,
+  "max_position_embeddings": 4096,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 1,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.0.dev0",
+  "use_cache": true,
+  "vocab_size": 32000
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "transformers_version": "4.57.0.dev0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db6eefbdc543761743970f19ef7a504ba5be8e7980fb2f30dfa6a5f44f9259c7
+size 214058792

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "</s>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,42 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": null,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 2048,
+  "pad_token": "</s>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

train_run1.py

@@ -0,0 +1,390 @@
+#!/usr/bin/env python3
+"""
+Nanomind pretraining script for decoder-only causal LM on JSONL.gz data.
+
+- Expects input file with one JSON object per line containing a `text` field.
+- Streams, tokenizes, and packs sequences to a fixed length for efficient training.
+- Uses a small LLaMA-style config by default (RMSNorm + SwiGLU + RoPE, MQA).
+
+Usage example:
+  python /workspace/nanomind/train.py \
+    --data_path /workspace/nanomind_data/pretrain_1m.jsonl.gz \
+    --out_dir /workspace/nanomind_runs/run1 \
+    --tokenizer_name hf-internal-testing/llama-tokenizer \
+    --seq_len 4096 --global_batch_size 256 \
+    --lr 1e-3 --warmup_steps 2000 --max_steps 50000 --bf16
+"""
+
+import os
+import io
+import gc
+import gzip
+import json
+import math
+import time
+import random
+import argparse
+from pathlib import Path
+from typing import Iterator, List, Dict, Optional
+
+import torch
+from torch import nn
+from torch.utils.data import IterableDataset, DataLoader
+
+from transformers import (
+    AutoTokenizer,
+    LlamaConfig,
+    LlamaForCausalLM,
+    get_cosine_schedule_with_warmup,
+)
+
+
+class JsonlPackedDataset(IterableDataset):
+    """
+    Streams a JSONL(.gz) file of objects with a `text` field, tokenizes, and
+    packs tokens into fixed-length blocks of `seq_len`.
+    """
+
+    def __init__(
+        self,
+        data_path: str,
+        tokenizer,
+        seq_len: int,
+        shuffle_lines: bool = False,
+        add_bos_eos: bool = True,
+        repeat: bool = True,
+        buffer_tokens_limit: int = 4_000_000,
+    ) -> None:
+        super().__init__()
+        self.data_path = str(data_path)
+        self.tokenizer = tokenizer
+        self.seq_len = int(seq_len)
+        self.shuffle_lines = bool(shuffle_lines)
+        self.add_bos_eos = bool(add_bos_eos)
+        self.repeat = bool(repeat)
+        self.buffer_tokens_limit = int(buffer_tokens_limit)
+
+        # pack buffers
+        self._token_buffer: List[int] = []
+
+    def _line_iter(self) -> Iterator[str]:
+        path = self.data_path
+        is_gz = path.endswith(".gz")
+        open_fn = gzip.open if is_gz else open
+        mode = "rt"
+        while True:
+            with open_fn(path, mode, encoding="utf-8") as f:
+                for line in f:
+                    yield line
+            if not self.repeat:
+                break
+
+    def _yield_blocks(self) -> Iterator[Dict[str, torch.Tensor]]:
+        bos_id = getattr(self.tokenizer, "bos_token_id", None)
+        eos_id = getattr(self.tokenizer, "eos_token_id", None)
+
+        # local references for speed
+        token_buffer = self._token_buffer
+        seq_len = self.seq_len
+
+        for raw_line in self._line_iter():
+            raw_line = raw_line.strip()
+            if not raw_line:
+                continue
+            try:
+                obj = json.loads(raw_line)
+            except json.JSONDecodeError:
+                continue
+            text = obj.get("text")
+            if not text or len(text) < 10:
+                continue
+
+            if self.add_bos_eos and bos_id is not None and eos_id is not None:
+                encoded = self.tokenizer.encode(
+                    text, add_special_tokens=False
+                )
+                # Guard against rare None returns
+                if not encoded:
+                    continue
+                token_buffer.append(bos_id)
+                token_buffer.extend(encoded)
+                token_buffer.append(eos_id)
+            else:
+                encoded = self.tokenizer.encode(text, add_special_tokens=True)
+                if not encoded:
+                    continue
+                token_buffer.extend(encoded)
+
+            # If buffer grows too large, drop tail to constrain RAM
+            if len(token_buffer) > self.buffer_tokens_limit:
+                del token_buffer[: len(token_buffer) - self.buffer_tokens_limit]
+
+            # Emit fixed-length blocks
+            while len(token_buffer) >= seq_len:
+                block = token_buffer[:seq_len]
+                del token_buffer[:seq_len]
+
+                input_ids = torch.tensor(block, dtype=torch.long)
+                attention_mask = torch.ones_like(input_ids)
+                # Causal LM uses labels equal to inputs
+                yield {
+                    "input_ids": input_ids,
+                    "attention_mask": attention_mask,
+                    "labels": input_ids.clone(),
+                }
+
+    def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
+        # Worker-specific shard: in IterableDataset DataLoader workers receive cloned objects.
+        # To keep it simple and deterministic, don't split lines per-worker; rely on global batching.
+        return self._yield_blocks()
+
+
+def build_model_and_tokenizer(
+    tokenizer_name: Optional[str],
+    tokenizer_dir: Optional[str],
+    model_name: Optional[str],
+    vocab_size_override: Optional[int],
+    hidden_size: int,
+    n_layers: int,
+    n_heads: int,
+    n_kv_heads: int,
+    rope_theta: float,
+    max_position_embeddings: int,
+) -> tuple:
+    # Tokenizer
+    if tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True)
+    elif tokenizer_dir:
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir, use_fast=True)
+    else:
+        raise ValueError("Provide --tokenizer_name or --tokenizer_dir")
+
+    # Ensure pad token for batching; map to eos if missing (common for causal LMs)
+    if tokenizer.pad_token_id is None:
+        if tokenizer.eos_token_id is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+        else:
+            # Fallback: add a [PAD] token
+            tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+    vocab_size = vocab_size_override or len(tokenizer)
+
+    # Model
+    if model_name:
+        model = LlamaForCausalLM.from_pretrained(model_name)
+        # Resize embeddings if tokenizer changed
+        if model.get_input_embeddings().weight.shape[0] != vocab_size:
+            model.resize_token_embeddings(vocab_size)
+    else:
+        config = LlamaConfig(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,  # d_model
+            intermediate_size=int(hidden_size * 2.2),  # SwiGLU widen 2.0–2.5
+            num_hidden_layers=n_layers,
+            num_attention_heads=n_heads,
+            num_key_value_heads=n_kv_heads,
+            rms_norm_eps=1e-5,
+            rope_theta=rope_theta,
+            max_position_embeddings=max_position_embeddings,
+            tie_word_embeddings=True,
+        )
+        model = LlamaForCausalLM(config)
+
+    return model, tokenizer
+
+
+def get_dataloader(
+    data_path: str,
+    tokenizer,
+    seq_len: int,
+    micro_batch_size: int,
+    num_workers: int,
+) -> DataLoader:
+    dataset = JsonlPackedDataset(
+        data_path=data_path,
+        tokenizer=tokenizer,
+        seq_len=seq_len,
+        shuffle_lines=False,
+        add_bos_eos=True,
+        repeat=True,
+    )
+    return DataLoader(
+        dataset,
+        batch_size=micro_batch_size,
+        num_workers=num_workers,
+        pin_memory=True,
+        drop_last=True,
+        collate_fn=_collate_batch,
+    )
+
+
+def _collate_batch(features: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    # All are fixed-length; just stack
+    input_ids = torch.stack([f["input_ids"] for f in features], dim=0)
+    attention_mask = torch.stack([f["attention_mask"] for f in features], dim=0)
+    labels = torch.stack([f["labels"] for f in features], dim=0)
+    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
+
+
+def parse_args() -> argparse.Namespace:
+    ap = argparse.ArgumentParser()
+    # Data
+    ap.add_argument("--data_path", required=True, help="Path to JSONL(.gz) with {text}")
+    ap.add_argument("--seq_len", type=int, default=4096)
+    ap.add_argument("--num_workers", type=int, default=2)
+
+    # Tokenizer & Model
+    ap.add_argument("--tokenizer_name", default=None, help="HF tokenizer name")
+    ap.add_argument("--tokenizer_dir", default=None, help="Local dir of HF tokenizer")
+    ap.add_argument("--model_name", default=None, help="HF model name to continue from (CPT)")
+    ap.add_argument("--vocab_size_override", type=int, default=None)
+
+    # Small LLaMA-like config (used when --model_name not provided)
+    ap.add_argument("--hidden_size", type=int, default=768)
+    ap.add_argument("--n_layers", type=int, default=24)
+    ap.add_argument("--n_heads", type=int, default=12)
+    ap.add_argument("--n_kv_heads", type=int, default=1)
+    ap.add_argument("--rope_theta", type=float, default=1e6)
+    ap.add_argument("--max_position_embeddings", type=int, default=4096)
+
+    # Training
+    ap.add_argument("--out_dir", required=True)
+    ap.add_argument("--global_batch_size", type=int, default=256)
+    ap.add_argument("--micro_batch_size", type=int, default=None, help="Per-step batch size before grad accumulation")
+    ap.add_argument("--lr", type=float, default=1e-3)
+    ap.add_argument("--weight_decay", type=float, default=0.05)
+    ap.add_argument("--warmup_steps", type=int, default=2000)
+    ap.add_argument("--max_steps", type=int, default=50_000)
+    ap.add_argument("--save_every", type=int, default=2000)
+    ap.add_argument("--clip_grad", type=float, default=1.0)
+    ap.add_argument("--bf16", action="store_true")
+    ap.add_argument("--seed", type=int, default=42)
+
+    return ap.parse_args()
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    model, tokenizer = build_model_and_tokenizer(
+        tokenizer_name=args.tokenizer_name,
+        tokenizer_dir=args.tokenizer_dir,
+        model_name=args.model_name,
+        vocab_size_override=args.vocab_size_override,
+        hidden_size=args.hidden_size,
+        n_layers=args.n_layers,
+        n_heads=args.n_heads,
+        n_kv_heads=args.n_kv_heads,
+        rope_theta=args.rope_theta,
+        max_position_embeddings=args.max_position_embeddings,
+    )
+
+    model = model.to(device)
+
+    # Data
+    micro_bs = args.micro_batch_size or min( max(1, args.global_batch_size // 8), args.global_batch_size)
+    grad_accum = max(1, args.global_batch_size // micro_bs)
+    train_loader = get_dataloader(
+        data_path=args.data_path,
+        tokenizer=tokenizer,
+        seq_len=args.seq_len,
+        micro_batch_size=micro_bs,
+        num_workers=args.num_workers,
+    )
+
+    # Optimizer & Scheduler
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95))
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=args.warmup_steps,
+        num_training_steps=args.max_steps,
+    )
+
+    scaler = None
+    use_bf16 = args.bf16 and torch.cuda.is_available()
+    autocast_dtype = torch.bfloat16 if use_bf16 else torch.float16
+
+    model.train()
+    step = 0
+    running_loss = 0.0
+    tokens_per_step = args.global_batch_size * args.seq_len
+    last_log = time.time()
+
+    # Simple training loop over streaming dataloader
+    data_iter = iter(train_loader)
+    while step < args.max_steps:
+        optimizer.zero_grad(set_to_none=True)
+        for micro_step in range(grad_accum):
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                data_iter = iter(train_loader)
+                batch = next(data_iter)
+
+            input_ids = batch["input_ids"].to(device, non_blocking=True)
+            attention_mask = batch["attention_mask"].to(device, non_blocking=True)
+            labels = batch["labels"].to(device, non_blocking=True)
+
+            with torch.autocast(device_type="cuda", dtype=autocast_dtype, enabled=use_bf16):
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+                loss = outputs.loss / grad_accum
+
+            loss.backward()
+            running_loss += loss.item()
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=args.clip_grad)
+        optimizer.step()
+        scheduler.step()
+        step += 1
+
+        # Logging
+        if step % 10 == 0:
+            now = time.time()
+            dt = now - last_log
+            last_log = now
+            avg_loss = running_loss / 10
+            running_loss = 0.0
+            ppl = math.exp(avg_loss) if avg_loss < 30 else float("inf")
+            tokens_sec = tokens_per_step / dt if dt > 0 else 0.0
+            print(
+                f"step {step:6d} | loss {avg_loss:.4f} | ppl {ppl:.2f} | tokens/s {tokens_sec:,.0f} | lr {scheduler.get_last_lr()[0]:.2e}",
+                flush=True,
+            )
+
+        # Checkpointing
+        if step % args.save_every == 0 or step == args.max_steps:
+            ckpt_dir = out_dir / f"step_{step:06d}"
+            ckpt_dir.mkdir(parents=True, exist_ok=True)
+            model.save_pretrained(ckpt_dir)
+            tokenizer.save_pretrained(ckpt_dir)
+
+        # Small memory hygiene
+        if step % 100 == 0:
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    # Final save
+    model.save_pretrained(out_dir / "final")
+    tokenizer.save_pretrained(out_dir / "final")
+
+
+if __name__ == "__main__":
+    main()
+
+