scripts

qwen2.5-7b_c4_1B/scripts/bits_llm.py

@@ -0,0 +1,108 @@
+"""
+Bits-for-LLM encoder (extended): token id -> fixed random n-bit code -> encoder -> hidden.
+Backward compatible: defaults (enc_type='mlp', ffn_depth=2, out_scale=0) reproduce the
+original 2-layer MLP exactly. New encoder-only tuning knobs for the 7B sweep:
+  - ffn_depth : number of Linear layers in the MLP (2 = original, 3 = deeper)
+  - out_scale : if >0, rescale the final projection so output std ~= out_scale (init-scale match)
+  - enc_type='tfenc' : small Transformer encoder over the bit-code (TF-Enc-H4-FFN512 replica)
+"""
+import torch, torch.nn as nn
+
+
+class BitsEmbedding(nn.Module):
+    def __init__(self, vocab, hidden, nbits=64, W=256, seed=0, dtype=torch.bfloat16,
+                 enc_type='mlp', ffn_depth=2, out_scale=0.0):
+        super().__init__()
+        g = torch.Generator().manual_seed(seed)
+        # plain attribute (NOT a buffer) so FSDP never shards it
+        self._codes = torch.randint(0, 2, (vocab, nbits), generator=g).to(dtype)
+        self.vocab, self.hidden, self.nbits, self.W = vocab, hidden, nbits, W
+        self.enc_type, self.ffn_depth = enc_type, ffn_depth
+        self.padding_idx = None
+        if enc_type == 'tfenc':
+            self.tok_bits = 16
+            assert nbits % self.tok_bits == 0, "nbits must be divisible by 16 for tfenc"
+            self.seq = nbits // self.tok_bits
+            self.in_proj = nn.Linear(self.tok_bits, W)
+            self.pos = nn.Parameter(torch.zeros(self.seq, W))
+            layer = nn.TransformerEncoderLayer(d_model=W, nhead=4, dim_feedforward=512,
+                                               activation='gelu', batch_first=True, norm_first=True)
+            self.tfenc = nn.TransformerEncoder(layer, num_layers=2)
+            self.out_proj = nn.Linear(W, hidden)
+            self.to(dtype)
+        else:  # mlp with ffn_depth linear layers (>=2)
+            layers = [nn.Linear(nbits, W), nn.GELU()]
+            for _ in range(max(0, ffn_depth - 2)):
+                layers += [nn.Linear(W, W), nn.GELU()]
+            layers += [nn.Linear(W, hidden)]
+            self.ffn = nn.Sequential(*layers).to(dtype)
+        if out_scale and out_scale > 0:
+            self._match_out_scale(out_scale)
+
+    def _final(self):
+        return self.out_proj if self.enc_type == 'tfenc' else self.ffn[-1]
+
+    def _match_out_scale(self, target_std):
+        # empirically rescale the final (linear) layer so the OUTPUT embedding std ~= target_std,
+        # matching the standard embedding's init scale (~initializer_range). Output is linear in
+        # the final layer for both mlp (ffn[-1]) and tfenc (out_proj), so scaling it scales output.
+        with torch.no_grad():
+            idx = torch.randint(0, self.vocab, (4096,))
+            cur = self.forward(idx).float().std().item()
+            if cur > 1e-8:
+                k = target_std / cur
+                last = self._final()
+                last.weight.mul_(k)
+                if last.bias is not None:
+                    last.bias.mul_(k)
+
+    def forward(self, ids):
+        if self._codes.device != ids.device:
+            self._codes = self._codes.to(ids.device)
+        codes = self._codes[ids]
+        if self.enc_type == 'tfenc':
+            *lead, n = codes.shape
+            x = codes.reshape(-1, self.seq, self.tok_bits)
+            x = self.in_proj(x) + self.pos
+            x = self.tfenc(x).mean(dim=1)
+            return self.out_proj(x).reshape(*lead, self.hidden)
+        return self.ffn(codes)
+
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+
+def swap_in_bits(model, nbits, W, seed, enc_type='mlp', ffn_depth=2, out_scale=0.0):
+    cfg = model.config
+    if getattr(cfg, "tie_word_embeddings", False):
+        out = model.get_output_embeddings()
+        if out is not None:
+            new = nn.Linear(out.weight.shape[1], out.weight.shape[0], bias=False)
+            new.weight = nn.Parameter(out.weight.detach().clone())
+            model.set_output_embeddings(new); cfg.tie_word_embeddings = False
+    dtype = next(model.parameters()).dtype
+    bits = BitsEmbedding(cfg.vocab_size, cfg.hidden_size, nbits, W, seed, dtype,
+                         enc_type=enc_type, ffn_depth=ffn_depth, out_scale=out_scale)
+    dev = model.get_input_embeddings().weight.device
+    model.set_input_embeddings(bits.to(dev))
+    orig = cfg.vocab_size * cfg.hidden_size
+    print(f"[Bits enc={enc_type} depth={ffn_depth} W={W} nbits={nbits} scale={out_scale}] "
+          f"table {orig/1e9:.3f}B -> encoder {bits.n_params()/1e6:.2f}M "
+          f"= {orig/max(bits.n_params(),1):.0f}x compression", flush=True)
+    return bits
+
+
+def get_dataset(tokenizer, max_len, n_tokens, split_stream):
+    from datasets import load_dataset
+    ds = load_dataset("allenai/c4", "en", split=split_stream, streaming=True)
+    def gen():
+        buf = []
+        while True:
+            for ex in ds:
+                ids = tokenizer(ex["text"], add_special_tokens=False)["input_ids"]
+                buf.extend(ids + [tokenizer.eos_token_id])
+                while len(buf) >= max_len:
+                    yield {"input_ids": buf[:max_len], "labels": buf[:max_len]}
+                    buf = buf[max_len:]
+    from datasets import IterableDataset
+    return IterableDataset.from_generator(gen)

qwen2.5-7b_c4_1B/scripts/bits_lm_traj.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+"""EXACT original from-scratch recipe (bits_lm_scratch.py) + intermediate checkpoints every 100M
+tokens, so we get the full 100M->500M trajectory for both arms. Training dynamics are identical to
+the original (only checkpoint saving is added).
+Camera-ready additions (2026-11-05): transformers 5 warmup_steps (= ceil(0.02*max_steps), what
+warmup_ratio=0.02 computed before), multi-GPU token accounting via WORLD_SIZE, optional FSDP
+sharding (same optimizer step, states sharded across ranks), keep all checkpoints."""
+import argparse, sys, math, os, torch
+sys.path.insert(0, '/root')
+from bits_llm import swap_in_bits, get_dataset
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, TrainingArguments, Trainer, TrainerCallback
+import subprocess
+
+
+class Bf16Save(TrainerCallback):
+    """Disk only: after each save, rewrite the fp32 safetensors of that checkpoint to bf16 in a detached
+    process (30G -> 15G per 7B checkpoint). Training dynamics untouched."""
+    def on_save(self, args, state, control, **kw):
+        if state.is_world_process_zero:
+            d = os.path.join(args.output_dir, f'checkpoint-{state.global_step}')
+            subprocess.Popen([sys.executable, '/root/to_bf16.py', d], stdout=open(d + '/bf16.log', 'w'), stderr=subprocess.STDOUT)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--arm', required=True, choices=['bits', 'standard'])
+    ap.add_argument('--model', default='Qwen/Qwen2.5-7B')
+    ap.add_argument('--tokens', type=int, default=500_000_000)
+    ap.add_argument('--max_len', type=int, default=512)
+    ap.add_argument('--bs', type=int, default=8)
+    ap.add_argument('--grad_accum', type=int, default=8)
+    ap.add_argument('--lr', type=float, default=3e-4)
+    ap.add_argument('--nbits', type=int, default=256)
+    ap.add_argument('--W', type=int, default=512)
+    ap.add_argument('--seed', type=int, default=42)
+    ap.add_argument('--out', required=True)
+    ap.add_argument('--fsdp', action='store_true')
+    ap.add_argument('--save_tokens', type=int, default=100_000_000)
+    a = ap.parse_args()
+    torch.manual_seed(a.seed)
+    tok = AutoTokenizer.from_pretrained(a.model)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    cfg = AutoConfig.from_pretrained(a.model)
+    model = AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.bfloat16, attn_implementation='sdpa')
+    if not a.fsdp:
+        model = model.cuda()
+    if a.arm == 'bits':
+        swap_in_bits(model, a.nbits, a.W, a.seed)
+        if not a.fsdp:
+            model = model.cuda()
+    nparam = sum(p.numel() for p in model.parameters())
+    emb = sum(p.numel() for p in model.get_input_embeddings().parameters())
+    print(f"[{a.arm} s{a.seed}] total={nparam/1e6:.0f}M  embed={emb/1e6:.2f}M", flush=True)
+    train_ds = get_dataset(tok, a.max_len, a.tokens, 'train')
+    world = int(os.environ.get('WORLD_SIZE', 1))
+    tps = a.max_len * a.bs * a.grad_accum * world   # tokens per optimizer step across all ranks
+    max_steps = max(1, a.tokens // tps)
+    save_every = max(1, a.save_tokens // tps)       # checkpoint every 100M tokens
+    extra = {}
+    if a.fsdp:
+        extra = dict(fsdp='full_shard auto_wrap',
+                     fsdp_config={'transformer_layer_cls_to_wrap': ['Qwen2DecoderLayer']})
+    targs = TrainingArguments(
+        output_dir=a.out, per_device_train_batch_size=a.bs, gradient_accumulation_steps=a.grad_accum,
+        learning_rate=a.lr, max_steps=max_steps, logging_steps=10, bf16=True,
+        save_strategy='steps', save_steps=save_every, save_total_limit=None, save_only_model=True,
+        optim='adamw_torch', gradient_checkpointing=True, report_to=['tensorboard'], dataloader_num_workers=2,
+        warmup_steps=math.ceil(0.02 * max_steps), lr_scheduler_type='cosine', max_grad_norm=1.0, seed=a.seed,
+        **extra)
+    def collate(b):
+        return {'input_ids': torch.tensor([x['input_ids'] for x in b]),
+                'labels': torch.tensor([x['labels'] for x in b])}
+    tr = Trainer(model=model, args=targs, train_dataset=train_ds, data_collator=collate)
+    tr.add_callback(Bf16Save())
+    print(f"[{a.arm} s{a.seed}] world={world} tokens/step={tps} max_steps={max_steps} save_every={save_every}", flush=True)
+    tr.train()
+    tr.save_model(a.out + '/final')
+    final = tr.state.log_history[-1].get('loss', float('nan')) if tr.state.log_history else float('nan')
+    print(f"[{a.arm} s{a.seed}] max_steps={max_steps} save_every={save_every} DONE  final_loss~={final}", flush=True)
+
+
+if __name__ == '__main__':
+    main()

qwen2.5-7b_c4_1B/scripts/eval_ppl.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""Held-out perplexity eval for Bits vs Standard checkpoints (supports encoder variants).
+Same held-out set, same metric exp(mean per-token CE). For Bits variants, pass the SAME
+--enc_type/--ffn_depth/--W/--nbits used at training so the module structure matches for loading."""
+import argparse, glob, os, math, sys, torch
+os.environ['HF_HUB_OFFLINE'] = '1'
+os.environ['TRANSFORMERS_OFFLINE'] = '1'
+import torch.nn.functional as F
+sys.path.insert(0, '/root')
+from transformers import AutoConfig, AutoModelForCausalLM
+from safetensors.torch import load_file
+
+
+def build(a):
+    cfg = AutoConfig.from_pretrained(a.ckpt)
+    if a.arm == 'standard':
+        return AutoModelForCausalLM.from_pretrained(
+            a.ckpt, torch_dtype=torch.bfloat16, attn_implementation='sdpa').cuda().eval()
+    from bits_llm import swap_in_bits
+    model = AutoModelForCausalLM.from_config(cfg, torch_dtype=torch.bfloat16, attn_implementation='sdpa')
+    swap_in_bits(model, a.nbits, a.W, a.seed, enc_type=a.enc_type, ffn_depth=a.ffn_depth, out_scale=0.0)
+    sd = {}
+    for f in glob.glob(os.path.join(a.ckpt, '*.safetensors')):
+        sd.update(load_file(f))
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    enc_keys = [k for k in sd if 'embed_tokens' in k or '.ffn.' in k or 'in_proj' in k or 'out_proj' in k or 'tfenc' in k]
+    print(f"[load] tensors={len(sd)} enc_keys={len(enc_keys)} missing={len(missing)} unexpected={len(unexpected)}", flush=True)
+    if len(enc_keys) == 0:
+        print("[load][FATAL] no encoder weights in checkpoint", flush=True); sys.exit(3)
+    return model.cuda().eval()
+
+
+@torch.no_grad()
+def evaluate(model, seqs, bs):
+    tot_loss, tot_tok = 0.0, 0
+    for i in range(0, len(seqs), bs):
+        batch = torch.tensor(seqs[i:i+bs]).cuda()
+        logits = model(input_ids=batch).logits[:, :-1, :].float()
+        labels = batch[:, 1:]
+        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1), reduction='sum')
+        tot_loss += loss.item(); tot_tok += labels.numel()
+    ml = tot_loss / tot_tok
+    return ml, math.exp(ml)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument('--arm', required=True, choices=['bits', 'standard'])
+    ap.add_argument('--ckpt', required=True)
+    ap.add_argument('--data', default='/root/heldout_c4val_4000x512.pt')
+    ap.add_argument('--nbits', type=int, default=256)
+    ap.add_argument('--W', type=int, default=512)
+    ap.add_argument('--enc_type', default='mlp', choices=['mlp', 'tfenc'])
+    ap.add_argument('--ffn_depth', type=int, default=2)
+    ap.add_argument('--seed', type=int, default=42)
+    ap.add_argument('--n_seq', type=int, default=4000)
+    ap.add_argument('--bs', type=int, default=8)
+    a = ap.parse_args()
+    model = build(a)
+    seqs = torch.load(a.data)[:a.n_seq]
+    ml, ppl = evaluate(model, seqs, a.bs)
+    print(f"RESULT arm={a.arm} ckpt={a.ckpt} enc={a.enc_type} depth={a.ffn_depth} W={a.W} nbits={a.nbits} "
+          f"tokens={len(seqs)*len(seqs[0])} eval_loss={ml:.4f} ppl={ppl:.2f}", flush=True)
+
+
+if __name__ == '__main__':
+    main()

qwen2.5-7b_c4_1B/scripts/eval_seed.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+# usage: eval_seed.sh SEED ARM GPU  -> appends RESULT lines to /root/traj/s3_1B/eval_s${SEED}_${ARM}.log
+seed=$1; arm=$2; gpu=$3
+source /venv/main/bin/activate; cd /root
+dir=/root/traj/s3_1B/s${seed}_$([ $arm = bits ] && echo bits || echo std)
+out=/root/traj/s3_1B/eval_s${seed}_${arm}.log
+for c in $(ls -d $dir/checkpoint-* | sort -t- -k2 -n) $dir/final; do
+  grep -q "ckpt=$c " $out 2>/dev/null && continue
+  CUDA_VISIBLE_DEVICES=$gpu python eval_ppl.py --arm $arm --ckpt $c --seed $seed 2>&1 | grep -aE "RESULT|FATAL|Traceback|Error" >> $out
+done
+echo "EVAL_DONE seed=$seed arm=$arm" >> $out

qwen2.5-7b_c4_1B/scripts/heldout_c4val_4000x512.pt

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

qwen2.5-7b_c4_1B/scripts/hf_arm.py

@@ -0,0 +1,18 @@
+import os, sys, time
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"]="1"
+from huggingface_hub import HfApi
+arm=sys.argv[1]; tok=sys.stdin.readline().strip()
+api=HfApi(token=tok); repo="bit2vec/bitsseal"; base="/root/traj/s3_1B"
+for attempt in range(8):
+    try:
+        api.upload_folder(folder_path=f"{base}/{arm}", repo_id=repo, path_in_repo=f"qwen2.5-7b_c4_1B/{arm}",
+            allow_patterns=["checkpoint-*/*.safetensors","checkpoint-*/*.json","checkpoint-*/bf16.log"], commit_message=f"{arm}: all checkpoints")
+        print("ARM_UPLOADED", arm, time.strftime("%H:%M:%S"), flush=True); break
+    except Exception as e:
+        print("RETRY", arm, attempt, repr(e)[:300], flush=True); time.sleep(300)
+else:
+    print("GAVEUP", arm, flush=True)
+if arm=="s42_bits":
+    api.upload_folder(folder_path=base, repo_id=repo, path_in_repo="qwen2.5-7b_c4_1B/logs", allow_patterns=["*.log"], ignore_patterns=["**/checkpoint-*/**","**/runs/**"], commit_message="logs")
+    api.upload_folder(folder_path="/root", repo_id=repo, path_in_repo="qwen2.5-7b_c4_1B/scripts", allow_patterns=["bits_llm.py","bits_lm_traj.py","eval_ppl.py","eval_seed.sh","run_3seeds_1B.sh","to_bf16.py","hf_arm.py","heldout_c4val_4000x512.pt"], commit_message="scripts")
+    print("LOGS_SCRIPTS_DONE", flush=True)

qwen2.5-7b_c4_1B/scripts/run_3seeds_1B.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# 3 seed pairs at 1B tokens, both arms of a seed at once on 4 GPUs each, seeds back to back.
+cd /root
+export FSDP_MODE="full_shard auto_wrap"
+OUT=/root/traj/s3_1B; mkdir -p $OUT
+for seed in 43 44 42; do
+  echo "[$(date -u +%FT%TZ)] START seed $seed" >> $OUT/driver.log
+  CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node 4 --master_port 29501 bits_lm_traj.py --arm bits     --bs 16 --grad_accum 1 --tokens 1000000000 --seed $seed --fsdp --out $OUT/s${seed}_bits > $OUT/s${seed}_bits.log 2>&1 &
+  P1=$!
+  CUDA_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node 4 --master_port 29502 bits_lm_traj.py --arm standard --bs 16 --grad_accum 1 --tokens 1000000000 --seed $seed --fsdp --out $OUT/s${seed}_std  > $OUT/s${seed}_std.log 2>&1 &
+  P2=$!
+  wait $P1; R1=$?; wait $P2; R2=$?
+  echo "[$(date -u +%FT%TZ)] END seed $seed bits_rc=$R1 std_rc=$R2" >> $OUT/driver.log
+done
+echo "[$(date -u +%FT%TZ)] ALL_SEEDS_DONE" >> $OUT/driver.log

qwen2.5-7b_c4_1B/scripts/to_bf16.py

@@ -0,0 +1,11 @@
+import sys, os, glob, torch
+from safetensors.torch import load_file, save_file
+d = sys.argv[1]
+for p in sorted(glob.glob(os.path.join(d, '*.safetensors'))):
+    sd = load_file(p)
+    n = sum(v.numel() for v in sd.values())
+    sd = {k: (v.to(torch.bfloat16) if v.is_floating_point() else v) for k, v in sd.items()}
+    save_file(sd, p + '.tmp', metadata={'format': 'pt'})
+    os.replace(p + '.tmp', p)
+    print('bf16', p, n, flush=True)
+print('DONE', flush=True)