@tuned-tensor/cli 0.13.1-beta.1 → 0.13.1-beta.2

package/training/foundation/src/pretrain.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from tokenizers import Tokenizer
+from torch.nn import functional as F
+
+from common import load_config, require_cuda, write_json
+from data import corpus_hash, decoded_byte_count, example_pairs, smoke_corpus, window_tokens
+from model import FoundationGPT, bits_per_byte, model_config_from_depth, param_count, save_model
+
+
+def main(argv: list[str] | None = None) -> None:
+    config = load_config(argv)
+    require_cuda("pretrain")
+    device = torch.device("cuda")
+    tokenizer = Tokenizer.from_file(str(Path(config["tokenizer_dir"]) / "tokenizer.json"))
+    examples = example_pairs(config.get("examples") or [])
+    corpus = smoke_corpus(examples, int(config.get("max_chars") or 20_000), extra=str(config.get("system_prompt") or ""))
+    token_ids = tokenizer.encode(corpus).ids
+    sequence_length = int(config.get("sequence_length") or 64)
+    windows = window_tokens(token_ids, sequence_length)
+    depth = int(config.get("depth") or 2)
+    model_config = model_config_from_depth(depth, tokenizer.get_vocab_size(), sequence_length)
+    model = FoundationGPT(model_config).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
+    batch_size = max(1, int(config.get("batch_size") or 2))
+    steps = max(1, int(config.get("steps") or 2))
+    last_loss = 0.0
+    last_bpb = 0.0
+    model.train()
+    for step in range(steps):
+        batch = windows[step % len(windows): step % len(windows) + batch_size]
+        if len(batch) < batch_size:
+            batch = (batch + windows)[:batch_size]
+        tensor = torch.tensor(batch, dtype=torch.long, device=device)
+        inputs = tensor[:, :-1]
+        labels = tensor[:, 1:]
+        logits = model(inputs)
+        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1))
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+        last_loss = float(loss.item())
+        byte_count = decoded_byte_count(tokenizer, labels.detach().cpu().tolist())
+        last_bpb = bits_per_byte(loss.detach() * labels.numel(), max(byte_count, 1))
+
+    output_dir = Path(config["output_dir"])
+    save_model(model, output_dir)
+    write_json(output_dir / "metrics.json", {
+        "ok": True,
+        "steps": steps,
+        "loss": last_loss,
+        "bpb": last_bpb,
+        "parameters": param_count(model),
+        "depth": depth,
+        "width": model_config["width"],
+        "heads": model_config["heads"],
+        "corpus_sha256": corpus_hash(corpus),
+        "device": str(device),
+    })
+
+
+if __name__ == "__main__":
+    main()

package/training/foundation/src/rl.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from tokenizers import Tokenizer
+from torch.nn import functional as F
+
+from common import load_config, require_cuda, write_json
+from data import END, encode_ids, example_pairs, format_prompt, numeric_reward
+from model import FoundationGPT, load_model, param_count, save_model
+
+
+def bounded_rollout(
+    prompt_ids: list[int],
+    context_length: int,
+    requested_completion_tokens: int = 24,
+) -> tuple[list[int], int]:
+    """Fit the newest prompt tokens and a useful completion inside one context."""
+    if context_length < 2:
+        raise ValueError("RL requires a model context of at least two tokens.")
+    completion_tokens = min(requested_completion_tokens, context_length // 2)
+    prompt_tokens = context_length - completion_tokens
+    return (prompt_ids or [0])[-prompt_tokens:], completion_tokens
+
+
+def policy_gradient_loss(token_log_probs: torch.Tensor, reward: float) -> torch.Tensor:
+    """Use a fixed binary-reward baseline so failed samples are discouraged."""
+    advantage = reward - 0.5
+    return -(advantage * token_log_probs.mean())
+
+
+@torch.no_grad()
+def sample_rollout(
+    model: FoundationGPT,
+    prompt: torch.Tensor,
+    max_new_tokens: int,
+    *,
+    temperature: float = 1.0,
+    stop_token_id: int | None = None,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Sample one on-policy completion, bounded by the model context."""
+    if temperature <= 0:
+        raise ValueError("RL sampling temperature must be positive.")
+    model.eval()
+    tokens = prompt
+    limit = int(model.config["sequence_length"])
+    for _ in range(max_new_tokens):
+        logits = model(tokens[:, -limit:])[:, -1, :] / temperature
+        probabilities = torch.softmax(logits, dim=-1)
+        next_token = torch.multinomial(probabilities, num_samples=1, generator=generator)
+        tokens = torch.cat([tokens, next_token], dim=1)
+        if stop_token_id is not None and bool(torch.all(next_token == stop_token_id)):
+            break
+    return tokens
+
+
+def main(argv: list[str] | None = None) -> None:
+    config = load_config(argv)
+    require_cuda("rl")
+    device = torch.device("cuda")
+    tokenizer = Tokenizer.from_file(str(Path(config["tokenizer_dir"]) / "tokenizer.json"))
+    model = load_model(config["model_dir"], map_location=device)
+    examples = example_pairs(config.get("examples") or [])
+    if not examples:
+        raise SystemExit("Foundation RL needs chat examples with verifiable answers.")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    generator = torch.Generator(device=device)
+    generator.manual_seed(int(config.get("seed") or 0))
+    steps = max(1, int(config.get("steps") or 1))
+    last_reward = 0.0
+    last_loss = 0.0
+    for step in range(steps):
+        example = examples[step % len(examples)]
+        prompt_ids = encode_ids(tokenizer, format_prompt(str(config.get("system_prompt") or ""), example["input"]))
+        prompt_ids, completion_tokens = bounded_rollout(
+            prompt_ids,
+            int(model.config["sequence_length"]),
+        )
+        prompt = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+        sampled = sample_rollout(
+            model,
+            prompt,
+            max_new_tokens=completion_tokens,
+            stop_token_id=tokenizer.token_to_id(END),
+            generator=generator,
+        )
+        completion_ids = sampled[0, len(prompt_ids):].tolist()
+        text = tokenizer.decode(completion_ids).split(END)[0]
+        reward = numeric_reward(text, example["output"])
+        model.train()
+        logits = model(sampled[:, :-1])
+        completion = sampled[:, len(prompt_ids):]
+        if completion.size(1) == 0:
+            last_reward = reward
+            continue
+        step_logits = logits[:, len(prompt_ids) - 1: len(prompt_ids) - 1 + completion.size(1)]
+        log_probs = F.log_softmax(step_logits, dim=-1)
+        token_log_probs = log_probs.gather(-1, completion.unsqueeze(-1)).squeeze(-1)
+        loss = policy_gradient_loss(token_log_probs, reward)
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        optimizer.step()
+        last_reward = reward
+        last_loss = float(loss.item())
+
+    output_dir = Path(config["output_dir"])
+    save_model(model, output_dir)
+    write_json(output_dir / "metrics.json", {
+        "ok": True,
+        "steps": steps,
+        "reward": last_reward,
+        "loss": last_loss,
+        "parameters": param_count(model),
+        "device": str(device),
+    })
+
+
+if __name__ == "__main__":
+    main()

package/training/foundation/src/train_tokenizer.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from common import ensure_private_directory, load_config, make_private_file, write_json
+from data import corpus_hash, example_pairs, smoke_corpus, train_tokenizer
+
+
+def main(argv: list[str] | None = None) -> None:
+    config = load_config(argv)
+    output_dir = ensure_private_directory(config["output_dir"])
+    examples = example_pairs(config.get("examples") or [])
+    corpus = smoke_corpus(examples, int(config.get("max_chars") or 20_000), extra=str(config.get("system_prompt") or ""))
+    tokenizer = train_tokenizer(corpus, int(config.get("vocab_size") or 256))
+    tokenizer_path = output_dir / "tokenizer.json"
+    tokenizer.save(str(tokenizer_path))
+    make_private_file(tokenizer_path)
+    write_json(output_dir / "metrics.json", {
+        "ok": True,
+        "vocab_size": tokenizer.get_vocab_size(),
+        "corpus_chars": len(corpus),
+        "corpus_sha256": corpus_hash(corpus),
+        "tokenizer_path": str(tokenizer_path),
+    })
+
+
+if __name__ == "__main__":
+    main()