| """Load released Safetensors checkpoints and generate text with the native model.""" |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
|
|
| from model import Config, GPT |
|
|
|
|
| def load_pretrained(model_path, *, subfolder=None, device="cuda", revision=None, |
| attention_backend="sdpa", moe_backend=None, mlp_backend="torch"): |
| """Return (GPT, SentencePieceProcessor) from a local folder or HF repository. |
| |
| Parameters remain FP32. Use BF16 autocast on CUDA, as in the paper. The |
| default SDPA/torch activation kernels do not require FlashAttention. |
| """ |
| from safetensors.torch import load_file |
| import sentencepiece as spm |
|
|
| if subfolder and (Path(subfolder).is_absolute() or ".." in Path(subfolder).parts): |
| raise ValueError("subfolder must be a relative path within the repository") |
| root = Path(model_path).expanduser() |
| if not root.is_dir(): |
| from huggingface_hub import snapshot_download |
| prefix = subfolder.rstrip("/") + "/" if subfolder else "" |
| root = Path(snapshot_download(str(model_path), revision=revision, |
| allow_patterns=[prefix + "*"])) |
| directory = root / subfolder if subfolder else root |
| fields = json.loads((directory / "model_config.json").read_text()) |
| target = torch.device(device) |
| fields.update(attention_backend=attention_backend, |
| moe_backend=moe_backend or ("grouped_mm" if target.type == "cuda" else "torch"), |
| mlp_backend=mlp_backend) |
| with torch.device("meta"): |
| model = GPT(Config(**fields)) |
| expected = {key: tuple(value.shape) for key, value in model.state_dict().items()} |
| index = json.loads((directory / "model.safetensors.index.json").read_text()) |
| weights = index["weight_map"] |
| if set(weights) != set(expected): |
| raise ValueError("Checkpoint tensor names differ from the configured model") |
| seen = set() |
| for name in sorted(set(weights.values())): |
| if Path(name).name != name or not name.endswith(".safetensors"): |
| raise ValueError("Invalid checkpoint shard filename") |
| state = load_file(str(directory / name), device="cpu") |
| declared = {key for key, shard in weights.items() if shard == name} |
| if set(state) != declared or seen.intersection(state): |
| raise ValueError("Missing, duplicated, or unexpected tensors in checkpoint shard") |
| for key, value in state.items(): |
| if tuple(value.shape) != expected[key] or value.dtype != torch.float32: |
| raise ValueError(f"Unexpected shape or precision for {key}") |
| model.load_state_dict(state, strict=False, assign=True) |
| seen.update(state) |
| del state |
| if seen != set(expected) or any(p.is_meta for p in model.parameters()): |
| raise ValueError("Checkpoint did not populate all model parameters") |
| model = model.to(target).eval() |
| tokenizer = spm.SentencePieceProcessor(model_file=str(directory / "tokenizer/tokenizer.model")) |
| if tokenizer.vocab_size() != model.config.vocab_size: |
| raise ValueError("Tokenizer vocabulary differs from the model") |
| return model, tokenizer |
|
|
|
|
| @torch.inference_mode() |
| def generate(model, token_ids, *, max_new_tokens=64, eos_token_id=2): |
| """Greedy continuation with each loop visit's own incremental KV cache.""" |
| if token_ids.ndim != 2 or token_ids.size(0) != 1 or token_ids.size(1) == 0: |
| raise ValueError("Expected a nonempty prompt with batch size one") |
| if max_new_tokens <= 0 or token_ids.size(1) + max_new_tokens > model.config.block_size: |
| raise ValueError("Prompt plus generated tokens must fit the model's context length") |
| tokens = token_ids |
| start = 0 |
| model.reset_cache() |
| try: |
| for _ in range(max_new_tokens): |
| segment = tokens[:, start:] |
| positions = torch.arange(start, tokens.size(1), device=tokens.device) |
| logits = model(segment, input_pos=positions) |
| next_token = logits[:, -1].argmax(-1, keepdim=True) |
| start = tokens.size(1) |
| tokens = torch.cat((tokens, next_token), dim=1) |
| if next_token.item() == eos_token_id: |
| break |
| return tokens |
| finally: |
| model.reset_cache() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model", default="yhx12/LatentLoop", help="HF repository ID or local folder") |
| parser.add_argument("--subfolder", default="1.3B-40B", help="1.3B-40B, 3B-70B, or 7B-140B") |
| parser.add_argument("--revision", help="Optional immutable HF commit") |
| parser.add_argument("--prompt", default="The capital of France is") |
| parser.add_argument("--max-new-tokens", type=int, default=64) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--attention-backend", choices=("sdpa", "flash"), default="sdpa") |
| parser.add_argument("--moe-backend", choices=("torch", "grouped_mm")) |
| parser.add_argument("--mlp-backend", choices=("torch", "flash"), default="torch") |
| args = parser.parse_args() |
| model, tokenizer = load_pretrained( |
| args.model, subfolder=args.subfolder, device=args.device, revision=args.revision, |
| attention_backend=args.attention_backend, moe_backend=args.moe_backend, |
| mlp_backend=args.mlp_backend) |
| prompt = tokenizer.encode(args.prompt, out_type=int) |
| tokens = torch.tensor([prompt], dtype=torch.long, device=args.device) |
| with torch.autocast("cuda", dtype=torch.bfloat16, enabled=torch.device(args.device).type == "cuda"): |
| output = generate(model, tokens, max_new_tokens=args.max_new_tokens, |
| eos_token_id=tokenizer.eos_id()) |
| print(tokenizer.decode(output[0].tolist())) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
