diff --git a/.gitignore b/.gitignore index 4a87b23d..7f280bd5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,9 @@ .venv/ __pycache__/ *.pyc -rustbpe/target/ dev-ignore/ report.md -eval_bundle/ \ No newline at end of file +eval_bundle/ + +# Secrets +.env \ No newline at end of file diff --git a/README.md b/README.md index 5e1de1ed..5f75429e 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ This repo is a full-stack implementation of an LLM like ChatGPT in a single, cle ## Talk to it -To get a sense of the endpoint of this repo, you can currently find [nanochat d32](https://github.com/karpathy/nanochat/discussions/8) hosted on [nanochat.karpathy.ai](https://nanochat.karpathy.ai/). "d32" means that this model has 32 layers in the Transformer neural network. This model has 1.9 billion parameters, it was trained on 38 billion tokens by simply running the single script [run1000.sh](run1000.sh), and the total cost of training was ~$800 (about 33 hours training time on 8XH100 GPU node). While today this is enough to outperform GPT-2 of 2019, it falls dramatically short of modern Large Language Models like GPT-5. When talking to these micro models, you'll see that they make a lot of mistakes, they are a little bit naive and silly and they hallucinate a ton, a bit like children. It's kind of amusing. But what makes nanochat unique is that it is fully yours - fully configurable, tweakable, hackable, and trained by you from start to end. To train and talk to your own, we turn to... +To get a sense of the endpoint of this repo, you can currently find [nanochat d34](https://github.com/karpathy/nanochat/discussions/314) hosted on [nanochat.karpathy.ai](https://nanochat.karpathy.ai/). "d34" means that this model has 34 layers in the Transformer neural network. This model has 2.2 billion parameters, it was trained on 88 billion tokens by simply running the training script [run1000.sh](run1000.sh) with `--target_param_data_ratio=40` (2x longer than Chinchilla-optimal), and the total cost of training was ~$2,500 (about 100 hours training time on 8XH100 GPU node). While today this is enough to outperform GPT-2 of 2019, it falls dramatically short of modern Large Language Models like GPT-5. When talking to these micro models, you'll see that they make a lot of mistakes, they are a little bit naive and silly and they hallucinate a ton, a bit like children. It's kind of amusing. But what makes nanochat unique is that it is fully yours - fully configurable, tweakable, hackable, and trained by you from start to end. To train and talk to your own, we turn to... ## Quick start @@ -108,10 +108,10 @@ Additionally, to add new abilities to nanochat, see [Guide: counting r in strawb nanochat is designed to be short and sweet. One big advantage of this is that we can package up all of the files together and copy paste them to your favorite LLM to ask arbitrary questions. As an example, I like to package up the repo using the [files-to-prompt](https://github.com/simonw/files-to-prompt) utility like so: ```bash -files-to-prompt . -e py -e md -e rs -e html -e toml -e sh --ignore "*target*" --cxml > packaged.txt +files-to-prompt . -e py -e md -e html -e toml -e sh --cxml > packaged.txt ``` -This includes all py, rs, html, toml, sh files, excludes the `rustbpe/target` folder, and chooses the cxml output format. Everything is written to the `packaged.txt` file, which atm measures ~330KB (i.e. well below ~100K tokens for a state of the art LLM), and ~8K lines of code in 45 files. +This includes all py, html, toml, sh files and chooses the cxml output format. Everything is written to the `packaged.txt` file, which atm measures ~330KB (i.e. well below ~100K tokens for a state of the art LLM), and ~8K lines of code in 45 files. Alternatively, I recommend using [DeepWiki](https://deepwiki.com/karpathy/nanochat) from Devin/Cognition to ask questions of this repo. In the URL of this repo, simply change github.com to deepwiki.com, and you're off. @@ -120,7 +120,7 @@ Alternatively, I recommend using [DeepWiki](https://deepwiki.com/karpathy/nanoch I haven't invested too much here but some tests exist, especially for the tokenizer. Run e.g. as: ```bash -python -m pytest tests/test_rustbpe.py -v -s +python -m pytest tests/test_engine.py -v -s ``` ## File structure @@ -140,7 +140,6 @@ python -m pytest tests/test_rustbpe.py -v -s │ ├── adamw.py # Distributed AdamW optimizer │ ├── checkpoint_manager.py # Save/Load model checkpoints │ ├── common.py # Misc small utilities, quality of life -│ ├── configurator.py # A superior alternative to argparse │ ├── core_eval.py # Evaluates base model CORE score (DCLM paper) │ ├── dataloader.py # Tokenizing Distributed Data Loader │ ├── dataset.py # Download/read utils for pretraining data @@ -155,12 +154,6 @@ python -m pytest tests/test_rustbpe.py -v -s │ └── ui.html # HTML/CSS/JS for nanochat frontend ├── pyproject.toml ├── run1000.sh # Train the ~$800 nanochat d32 -├── rustbpe # Custom Rust BPE tokenizer trainer -│ ├── Cargo.lock -│ ├── Cargo.toml -│ ├── README.md # see for why this even exists -│ └── src -│ └── lib.rs ├── scripts │ ├── base_eval.py # Base model: calculate CORE score │ ├── base_loss.py # Base model: calculate bits per byte, sample @@ -185,7 +178,6 @@ python -m pytest tests/test_rustbpe.py -v -s │ └── spellingbee.py # Task teaching model to spell/count letters ├── tests │ └── test_engine.py -│ └── test_rustbpe.py └── uv.lock ``` diff --git a/dev/gen_synthetic_data.py b/dev/gen_synthetic_data.py index 068824f4..c08c7e6f 100644 --- a/dev/gen_synthetic_data.py +++ b/dev/gen_synthetic_data.py @@ -24,8 +24,7 @@ prompt: manually generate any kind of entropy you can think of and include it in your prompts to maintain healthy and good diversity in the data. -NOTE: You need OpenRouter API key in a file called "openroutertoken.txt" in the root directory of the repo. - (obviously you can tune this arbitrarily to your liking) +NOTE: You need OPENROUTER_API_KEY set in .env or as an environment variable. NOTE: For more details see this discussion: https://github.com/karpathy/nanochat/discussions/139 """ import requests @@ -34,10 +33,12 @@ import os import copy import random from concurrent.futures import ThreadPoolExecutor, as_completed +from dotenv import load_dotenv from nanochat.common import get_base_dir -api_key = open("openroutertoken.txt", "r", encoding="utf-8").read().strip() +load_dotenv() +api_key = os.environ["OPENROUTER_API_KEY"] url = "https://openrouter.ai/api/v1/chat/completions" headers = { diff --git a/nanochat/adamw.py b/nanochat/adamw.py index db591de8..8816057e 100644 --- a/nanochat/adamw.py +++ b/nanochat/adamw.py @@ -27,6 +27,7 @@ class DistAdamW(torch.optim.Optimizer): for group in self.param_groups: params: list[Tensor] = group["params"] for base_i in range(len(params)): + assert params[base_i].shape[0] % world_size == 0, f"First dim of parameter shape {params[base_i].shape} must be divisible by world size {world_size}" grad = params[base_i].grad rank_size = grad.shape[0] // world_size grad_slice = torch.empty_like(grad[:rank_size]) diff --git a/nanochat/checkpoint_manager.py b/nanochat/checkpoint_manager.py index 63f257f2..99f260e4 100644 --- a/nanochat/checkpoint_manager.py +++ b/nanochat/checkpoint_manager.py @@ -34,6 +34,7 @@ def save_checkpoint(checkpoint_dir, step, model_data, optimizer_data, meta_data, logger.info(f"Saved metadata to: {meta_path}") # Note that optimizer state is sharded across ranks, so each rank must save its own. if optimizer_data is not None: + os.makedirs(checkpoint_dir, exist_ok=True) optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}_rank{rank:d}.pt") torch.save(optimizer_data, optimizer_path) logger.info(f"Saved optimizer state to: {optimizer_path}") @@ -93,11 +94,11 @@ def build_model(checkpoint_dir, step, device, phase): return model, tokenizer, meta_data -def find_largest_model(checkpoint_dir): +def find_largest_model(checkpoints_dir): # attempt to guess the model tag: take the biggest model available - model_tags = [f for f in os.listdir(checkpoint_dir) if os.path.isdir(os.path.join(checkpoint_dir, f))] + model_tags = [f for f in os.listdir(checkpoints_dir) if os.path.isdir(os.path.join(checkpoints_dir, f))] if not model_tags: - raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}") + raise FileNotFoundError(f"No checkpoints found in {checkpoints_dir}") # 1) normally all model tags are of the form d, try that first: candidates = [] for model_tag in model_tags: @@ -109,7 +110,7 @@ def find_largest_model(checkpoint_dir): candidates.sort(key=lambda x: x[0], reverse=True) return candidates[0][1] # 2) if that failed, take the most recently updated model: - model_tags.sort(key=lambda x: os.path.getmtime(os.path.join(checkpoint_dir, x)), reverse=True) + model_tags.sort(key=lambda x: os.path.getmtime(os.path.join(checkpoints_dir, x)), reverse=True) return model_tags[0] diff --git a/nanochat/common.py b/nanochat/common.py index 8f36f947..22559ce5 100644 --- a/nanochat/common.py +++ b/nanochat/common.py @@ -113,12 +113,24 @@ def print_banner(): """ print0(banner) -def is_ddp(): - # TODO is there a proper way - return int(os.environ.get('RANK', -1)) != -1 +def is_ddp_requested() -> bool: + """ + True if launched by torchrun (env present), even before init. + Used to decide whether we *should* initialize a PG. + """ + return all(k in os.environ for k in ("RANK", "LOCAL_RANK", "WORLD_SIZE")) + +def is_ddp_initialized() -> bool: + """ + True if torch.distributed is available and the process group is initialized. + Used at cleanup to avoid destroying a non-existent PG. + """ + return dist.is_available() and dist.is_initialized() def get_dist_info(): - if is_ddp(): + if is_ddp_requested(): + # We rely on torchrun's env to decide if we SHOULD init. + # (Initialization itself happens in compute init.) assert all(var in os.environ for var in ['RANK', 'LOCAL_RANK', 'WORLD_SIZE']) ddp_rank = int(os.environ['RANK']) ddp_local_rank = int(os.environ['LOCAL_RANK']) @@ -158,11 +170,11 @@ def compute_init(device_type="cuda"): # cuda|cpu|mps # Precision if device_type == "cuda": - torch.set_float32_matmul_precision("high") # uses tf32 instead of fp32 for matmuls + torch.backends.cuda.matmul.fp32_precision = "tf32" # uses tf32 instead of fp32 for matmuls # Distributed setup: Distributed Data Parallel (DDP), optional, and requires CUDA - ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() - if ddp and device_type == "cuda": + is_ddp_requested, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() + if is_ddp_requested and device_type == "cuda": device = torch.device("cuda", ddp_local_rank) torch.cuda.set_device(device) # make "cuda" default to this device dist.init_process_group(backend="nccl", device_id=device) @@ -173,11 +185,11 @@ def compute_init(device_type="cuda"): # cuda|cpu|mps if ddp_rank == 0: logger.info(f"Distributed world size: {ddp_world_size}") - return ddp, ddp_rank, ddp_local_rank, ddp_world_size, device + return is_ddp_requested, ddp_rank, ddp_local_rank, ddp_world_size, device def compute_cleanup(): """Companion function to compute_init, to clean things up before script exit""" - if is_ddp(): + if is_ddp_initialized(): dist.destroy_process_group() class DummyWandb: diff --git a/nanochat/configurator.py b/nanochat/configurator.py deleted file mode 100644 index ec1b76d1..00000000 --- a/nanochat/configurator.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -Poor Man's Configurator. Probably a terrible idea. Example usage: -$ python train.py config/override_file.py --batch_size=32 -this will first run config/override_file.py, then override batch_size to 32 - -The code in this file will be run as follows from e.g. train.py: ->>> exec(open('configurator.py').read()) - -So it's not a Python module, it's just shuttling this code away from train.py -The code in this script then overrides the globals() - -I know people are not going to love this, I just really dislike configuration -complexity and having to prepend config. to every single variable. If someone -comes up with a better simple Python solution I am all ears. -""" - -import os -import sys -from ast import literal_eval - -def print0(s="",**kwargs): - ddp_rank = int(os.environ.get('RANK', 0)) - if ddp_rank == 0: - print(s, **kwargs) - -for arg in sys.argv[1:]: - if '=' not in arg: - # assume it's the name of a config file - assert not arg.startswith('--') - config_file = arg - print0(f"Overriding config with {config_file}:") - with open(config_file) as f: - print0(f.read()) - exec(open(config_file).read()) - else: - # assume it's a --key=value argument - assert arg.startswith('--') - key, val = arg.split('=') - key = key[2:] - if key in globals(): - try: - # attempt to eval it it (e.g. if bool, number, or etc) - attempt = literal_eval(val) - except (SyntaxError, ValueError): - # if that goes wrong, just use the string - attempt = val - # ensure the types match ok - if globals()[key] is not None: - attempt_type = type(attempt) - default_type = type(globals()[key]) - assert attempt_type == default_type, f"Type mismatch: {attempt_type} != {default_type}" - # cross fingers - print0(f"Overriding: {key} = {attempt}") - globals()[key] = attempt - else: - raise ValueError(f"Unknown config key: {key}") diff --git a/nanochat/dataloader.py b/nanochat/dataloader.py index 32712984..41368021 100644 --- a/nanochat/dataloader.py +++ b/nanochat/dataloader.py @@ -26,20 +26,26 @@ def tokenizing_distributed_data_loader_with_state(B, T, split, tokenizer_threads ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() def document_batches(): parquet_paths = list_parquet_files() + assert len(parquet_paths) != 0, "No dataset parquet files found, did you run dataset.py?" parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:] resume_pq_idx = resume_state_dict["pq_idx"] if resume_state_dict is not None else 0 resume_rg_idx = resume_state_dict["rg_idx"] if resume_state_dict is not None else None + first_pass = True pq_idx = resume_pq_idx # we kick off parquet files at the resume index (or by default just 0) while True: # iterate infinitely (multi-epoch) + pq_idx = resume_pq_idx if first_pass else 0 while pq_idx < len(parquet_paths): # iterate over all parquet files filepath = parquet_paths[pq_idx] pf = pq.ParquetFile(filepath) # Start from resume point if resuming on same file, otherwise from DDP rank # I know this state resumption is a little bit tricky and a little bit hacky... sigh. - if resume_rg_idx is not None: + if first_pass and (resume_rg_idx is not None) and (pq_idx == resume_pq_idx): base_idx = resume_rg_idx // ddp_world_size # in units of ddp_world_size base_idx += 1 # advance by 1 so that we definitely don't repeat data after resuming rg_idx = base_idx * ddp_world_size + ddp_rank + if rg_idx >= pf.num_row_groups: + pq_idx += 1 + continue resume_rg_idx = None # set to None as we only want to do this a single time else: rg_idx = ddp_rank @@ -51,6 +57,7 @@ def tokenizing_distributed_data_loader_with_state(B, T, split, tokenizer_threads yield batch[i:i+tokenizer_batch_size], (pq_idx, rg_idx) rg_idx += ddp_world_size # advance to the next row group (in DDP) pq_idx += 1 # advance to the next parquet file + first_pass = False batches = document_batches() # Now emit batches of tokens. diff --git a/nanochat/engine.py b/nanochat/engine.py index d749d942..d4367fb1 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -19,7 +19,7 @@ from contextlib import contextmanager from collections import deque from nanochat.common import compute_init, autodetect_device_type from nanochat.checkpoint_manager import load_model -from contextlib import nullcontext +from contextlib import nullcontext # ----------------------------------------------------------------------------- # Calculator tool helpers @@ -107,17 +107,23 @@ class KVCache: # 1) validate the shapes assert self.kv_cache is None, "Cannot prefill a non-empty KV cache" assert other.kv_cache is not None, "Cannot prefill with a None KV cache" - for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)): - # ix 0: num_layers, 1: k/v, 2: batch_size, 3: num_heads, 4: seq_len, 5: head_dim - if ix in [0, 1, 3, 5]: - # num_layers, k/v, num_heads, head_dim must match - assert dim1 == dim2, f"Dim {ix} mismatch: {dim1} != {dim2}" - elif ix == 2: - # batch_size can be expanded - assert dim1 == dim2 or dim2 == 1, f"Batch dim mismatch: {dim1} != {dim2}" - elif ix == 4: - # seq_len: self must be longer than other - assert dim1 >= dim2, f"Seq len mismatch: {dim1} < {dim2}" + + # Extract dimensions explicitly + self_layers, self_kv, self_batch, self_heads, self_seq, self_head_dim = self.kv_shape + other_layers, other_kv, other_batch, other_heads, other_seq, other_head_dim = other.kv_shape + + # Validate dimensions + assert self_layers == other_layers, f"Layer count mismatch: {self_layers} != {other_layers}" + assert self_kv == other_kv, f"K/V dimension mismatch: {self_kv} != {other_kv}" + assert self_heads == other_heads, f"Head count mismatch: {self_heads} != {other_heads}" + assert self_head_dim == other_head_dim, f"Head dim mismatch: {self_head_dim} != {other_head_dim}" + + # Batch size can be expanded (other can be 1, self can be larger) + assert self_batch == other_batch or other_batch == 1, f"Batch size mismatch: {self_batch} vs {other_batch} (other must be 1 or equal)" + + # Sequence length: self must be longer than other + assert self_seq >= other_seq, f"Sequence length mismatch: {self_seq} < {other_seq}" + # 2) initialize the cache dtype, device = other.kv_cache.dtype, other.kv_cache.device self.kv_cache = torch.empty(self.kv_shape, dtype=dtype, device=device) @@ -143,11 +149,11 @@ class KVCache: self.kv_cache = torch.cat([self.kv_cache, additional_cache], dim=4).contiguous() self.kv_shape = self.kv_cache.shape # Insert k, v into the cache - self.kv_cache[layer_idx, 0, :, :, t0:t1] = k - self.kv_cache[layer_idx, 1, :, :, t0:t1] = v + self.kv_cache[layer_idx, 0, :, :, t0:t1, :] = k + self.kv_cache[layer_idx, 1, :, :, t0:t1, :] = v # Return the full cached keys/values up to current position (as a view) - key_view = self.kv_cache[layer_idx, 0, :, :, :t1] - value_view = self.kv_cache[layer_idx, 1, :, :, :t1] + key_view = self.kv_cache[layer_idx, 0, :, :, :t1, :] + value_view = self.kv_cache[layer_idx, 1, :, :, :t1, :] # Increment pos after the last layer of the Transformer processes if layer_idx == self.kv_cache.size(0) - 1: self.pos = t1 @@ -161,7 +167,7 @@ def sample_next_token(logits, rng, temperature=1.0, top_k=None): assert temperature >= 0.0, "temperature must be non-negative" if temperature == 0.0: return torch.argmax(logits, dim=-1, keepdim=True) - if top_k is not None: + if top_k is not None and top_k > 0: k = min(top_k, logits.size(-1)) vals, idx = torch.topk(logits, k, dim=-1) vals = vals / temperature @@ -217,9 +223,7 @@ class Engine: ) ids = torch.tensor([tokens], dtype=torch.long, device=device) logits = self.model.forward(ids, kv_cache=kv_cache_prefill) - logits = logits[:, -1, :] - next_ids = sample_next_token(logits, rng, temperature, top_k) # (B, 1) - sampled_tokens = next_ids[:, 0].tolist() + logits = logits[:, -1, :].expand(num_samples, -1) # (num_samples, vocab_size) # 2) Replicate the KV cache for each sample/row kv_length_hint = (len(tokens) + max_tokens) if max_tokens is not None else self.model.config.sequence_len @@ -236,7 +240,6 @@ class Engine: # 4) Main generation loop num_generated = 0 - first_iteration = True while True: # Stop condition: we've reached max tokens if max_tokens is not None and num_generated >= max_tokens: @@ -245,18 +248,9 @@ class Engine: if all(state.completed for state in row_states): break - # Get sampled tokens - either from prefill or from forward pass - if first_iteration: - # Use the tokens we already sampled from prefill - sampled_tokens = [sampled_tokens[0]] * num_samples # Broadcast first token to all rows - # TODO: we should sample a token for each row instead of broadcasting - first_iteration = False - else: - # Forward the model and get the next token for each row - logits = self.model.forward(ids, kv_cache=kv_cache_decode) # (B, T, vocab_size) - logits = logits[:, -1, :] # (B, vocab_size) at last time step - next_ids = sample_next_token(logits, rng, temperature, top_k) # (B, 1) - sampled_tokens = next_ids[:, 0].tolist() + # Sample the next token for each row + next_ids = sample_next_token(logits, rng, temperature, top_k) # (B, 1) + sampled_tokens = next_ids[:, 0].tolist() # Process each row: choose the next token, update state, optional tool use token_column = [] # contains the next token id along each row @@ -293,8 +287,10 @@ class Engine: # Yield the token column yield token_column, token_masks num_generated += 1 - # Prepare ids for next iteration + + # Prepare logits for next iteration ids = torch.tensor(token_column, dtype=torch.long, device=device).unsqueeze(1) + logits = self.model.forward(ids, kv_cache=kv_cache_decode)[:, -1, :] # (B, vocab_size) def generate_batch(self, tokens, num_samples=1, **kwargs): """ diff --git a/nanochat/gpt.py b/nanochat/gpt.py index 216343c2..e6027a96 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -41,12 +41,10 @@ def norm(x): def apply_rotary_emb(x, cos, sin): assert x.ndim == 4 # multihead attention d = x.shape[3] // 2 - x1, x2 = x[..., :d], x[..., d:] # split up last time into two halves + x1, x2 = x[..., :d], x[..., d:] # split up last dim into two halves y1 = x1 * cos + x2 * sin # rotate pairs of dims y2 = x1 * (-sin) + x2 * cos - out = torch.cat([y1, y2], 3) # re-assemble - out = out.to(x.dtype) # ensure input/output dtypes match - return out + return torch.cat([y1, y2], 3) class CausalSelfAttention(nn.Module): def __init__(self, config, layer_idx): @@ -98,8 +96,7 @@ class CausalSelfAttention(nn.Module): # First, each query attends to all the cached keys/values (i.e. full prefix) attn_mask = torch.zeros((Tq, Tk), dtype=torch.bool, device=q.device) # True = keep, False = mask prefix_len = Tk - Tq - if prefix_len > 0: # can't be negative but could be zero - attn_mask[:, :prefix_len] = True + attn_mask[:, :prefix_len] = True # Then, causal attention within this chunk attn_mask[:, prefix_len:] = torch.tril(torch.ones((Tq, Tq), dtype=torch.bool, device=q.device)) y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa) @@ -136,17 +133,22 @@ class Block(nn.Module): class GPT(nn.Module): - def __init__(self, config): + def __init__(self, config, pad_vocab_size_to=64): super().__init__() self.config = config + # For DDP, we want vocab_size divisible by world_size. Also, there are potential performance benefits, see: + # https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.resize_token_embeddings + padded_vocab_size = ((config.vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to) * pad_vocab_size_to + if padded_vocab_size != config.vocab_size: + print0(f"Padding vocab_size from {config.vocab_size} to {padded_vocab_size} to be divisible by {pad_vocab_size_to}") self.transformer = nn.ModuleDict({ - "wte": nn.Embedding(config.vocab_size, config.n_embd), + "wte": nn.Embedding(padded_vocab_size, config.n_embd), "h": nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]), }) - self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) - # To support meta device initialization, we init the rotary embeddings here, but it's fake + self.lm_head = nn.Linear(config.n_embd, padded_vocab_size, bias=False) + # To support meta device initialization, we init the rotary embeddings here, but it's just "fake" meta tensors only. # As for rotary_seq_len, these rotary embeddings are pretty small/cheap in memory, - # so let's just over-compute them, but assert fail if we ever reach that amount. + # so let's just over-compute them by 10X, but assert fail if we ever reach that amount. # In the future we can dynamically grow the cache, for now it's fine. self.rotary_seq_len = config.sequence_len * 10 # 10X over-compute should be enough, TODO make nicer? head_dim = config.n_embd // config.n_head @@ -155,35 +157,46 @@ class GPT(nn.Module): self.register_buffer("sin", sin, persistent=False) def init_weights(self): - self.apply(self._init_weights) - # zero out classifier weights - torch.nn.init.zeros_(self.lm_head.weight) - # zero out c_proj weights in all blocks + """ + Initialize the full model in this one function for maximum clarity. + + wte (embedding): normal, std=1.0 + lm_head: normal, std=0.001 + for each block: + attn.c_q: uniform, std=1/sqrt(n_embd) + attn.c_k: uniform, std=1/sqrt(n_embd) + attn.c_v: uniform, std=1/sqrt(n_embd) + attn.c_proj: zeros + mlp.c_fc: uniform, std=1/sqrt(n_embd) + mlp.c_proj: zeros + """ + + # Embedding and unembedding + torch.nn.init.normal_(self.transformer.wte.weight, mean=0.0, std=1.0) + torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001) + + # Transformer blocks: uniform init with bound = sqrt(3) * std (same standard deviation as normal) + n_embd = self.config.n_embd + s = 3**0.5 * n_embd**-0.5 # sqrt(3) multiplier makes sure Uniform achieves the same std as Normal for block in self.transformer.h: + torch.nn.init.uniform_(block.attn.c_q.weight, -s, s) # weights use Uniform to avoid outliers + torch.nn.init.uniform_(block.attn.c_k.weight, -s, s) + torch.nn.init.uniform_(block.attn.c_v.weight, -s, s) + torch.nn.init.zeros_(block.attn.c_proj.weight) # projections are zero + torch.nn.init.uniform_(block.mlp.c_fc.weight, -s, s) torch.nn.init.zeros_(block.mlp.c_proj.weight) - torch.nn.init.zeros_(block.attn.c_proj.weight) - # init the rotary embeddings + + # Rotary embeddings head_dim = self.config.n_embd // self.config.n_head cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim) self.cos, self.sin = cos, sin - # Cast the embeddings from fp32 to bf16: optim can tolerate it and it saves memory: both in the model and the activations + + # Cast token embeddings to bf16: optimizer can tolerate it and it saves memory if self.transformer.wte.weight.device.type == "cuda": self.transformer.wte.to(dtype=torch.bfloat16) - def _init_weights(self, module): - if isinstance(module, nn.Linear): - # https://arxiv.org/pdf/2310.17813 - fan_out = module.weight.size(0) - fan_in = module.weight.size(1) - std = 1.0 / math.sqrt(fan_in) * min(1.0, math.sqrt(fan_out / fan_in)) - torch.nn.init.normal_(module.weight, mean=0.0, std=std) - if module.bias is not None: - torch.nn.init.zeros_(module.bias) - elif isinstance(module, nn.Embedding): - torch.nn.init.normal_(module.weight, mean=0.0, std=1.0) - - # TODO: bump base theta more, e.g. 100K is more common more recently def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None): + # TODO: bump base theta more? e.g. 100K is more common more recently # autodetect the device from model embeddings if device is None: device = self.transformer.wte.weight.device @@ -221,8 +234,7 @@ class GPT(nn.Module): # Create the AdamW optimizer for the embedding and lm_head # Scale the LR for the AdamW parameters by ∝1/√dmodel (having tuned the LRs for 768 dim model) dmodel_lr_scale = (model_dim / 768) ** -0.5 - if rank == 0: - print(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}") + print0(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}") adam_groups = [ dict(params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale), dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale), @@ -260,19 +272,19 @@ class GPT(nn.Module): x = norm(x) # Forward the lm_head (compute logits) - softcap = 15 + softcap = 15 # smoothly cap the logits to the range [-softcap, softcap] + logits = self.lm_head(x) # (B, T, padded_vocab_size) <- very big tensor, large amount of memory + logits = logits[..., :self.config.vocab_size] # slice to remove padding + logits = logits.float() # switch to fp32 for logit softcap and loss computation + logits = softcap * torch.tanh(logits / softcap) # squash the logits + if targets is not None: - # training mode: compute and return the loss - # TODO: experiment with Liger Kernels / chunked cross-entropy etc. - logits = self.lm_head(x) - logits = softcap * torch.tanh(logits / softcap) # logits softcap - logits = logits.float() # use tf32/fp32 for logits + # training: given the targets, compute and return the loss + # TODO experiment with chunked cross-entropy? loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction) return loss else: - # inference mode: compute and return the logits - logits = self.lm_head(x) - logits = softcap * torch.tanh(logits / softcap) # logits softcap + # inference: just return the logits directly return logits @torch.inference_mode() diff --git a/nanochat/report.py b/nanochat/report.py index 0b0ebd73..1a31aa45 100644 --- a/nanochat/report.py +++ b/nanochat/report.py @@ -16,8 +16,11 @@ def run_command(cmd): """Run a shell command and return output, or None if it fails.""" try: result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=5) - if result.returncode == 0: + # Return stdout if we got output (even if some files in xargs failed) + if result.stdout.strip(): return result.stdout.strip() + if result.returncode == 0: + return "" return None except: return None @@ -160,12 +163,23 @@ Generated: {timestamp} """ - # bloat metrics: package all of the source code and assess its weight - packaged = run_command('files-to-prompt . -e py -e md -e rs -e html -e toml -e sh --ignore "*target*" --cxml') - num_chars = len(packaged) - num_lines = len(packaged.split('\n')) - num_files = len([x for x in packaged.split('\n') if x.startswith('')]) - num_tokens = num_chars // 4 # assume approximately 4 chars per token + # bloat metrics: count lines/chars in git-tracked source files only + extensions = ['py', 'md', 'rs', 'html', 'toml', 'sh'] + git_patterns = ' '.join(f"'*.{ext}'" for ext in extensions) + files_output = run_command(f"git ls-files -- {git_patterns}") + file_list = [f for f in (files_output or '').split('\n') if f] + num_files = len(file_list) + num_lines = 0 + num_chars = 0 + if num_files > 0: + wc_output = run_command(f"git ls-files -- {git_patterns} | xargs wc -lc 2>/dev/null") + if wc_output: + total_line = wc_output.strip().split('\n')[-1] + parts = total_line.split() + if len(parts) >= 2: + num_lines = int(parts[0]) + num_chars = int(parts[1]) + num_tokens = num_chars // 4 # assume approximately 4 chars per token # count dependencies via uv.lock uv_lock_lines = 0 diff --git a/nanochat/tokenizer.py b/nanochat/tokenizer.py index 880f8542..726fb2f8 100644 --- a/nanochat/tokenizer.py +++ b/nanochat/tokenizer.py @@ -122,7 +122,14 @@ class HuggingFaceTokenizer: return self.tokenizer.token_to_id(text) def get_bos_token_id(self): + # Different HuggingFace models use different BOS tokens and there is little consistency + # 1) attempt to find a <|bos|> token bos = self.encode_special("<|bos|>") + # 2) if that fails, attempt to find a <|endoftext|> token (e.g. GPT-2 models) + if bos is None: + bos = self.encode_special("<|endoftext|>") + # 3) if these fail, it's better to crash than to silently return None + assert bos is not None, "Failed to find BOS token in tokenizer" return bos def encode(self, text, *args, **kwargs): diff --git a/nanochat/ui.html b/nanochat/ui.html index 0f625d91..b85532a3 100644 --- a/nanochat/ui.html +++ b/nanochat/ui.html @@ -14,6 +14,11 @@ box-sizing: border-box; } + html, body{ + height: 100%; + margin: 0; + } + body { font-family: ui-sans-serif, -apple-system, system-ui, "Segoe UI", Helvetica, "Apple Color Emoji", Arial, sans-serif, "Segoe UI Emoji", "Segoe UI Symbol"; background-color: #ffffff; @@ -107,7 +112,6 @@ .message.assistant .message-content { background: transparent; border: none; - padding: 0.25rem 0; cursor: pointer; border-radius: 0.5rem; padding: 0.5rem; diff --git a/pyproject.toml b/pyproject.toml index 3d03c4bc..1f2234aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,30 +7,21 @@ requires-python = ">=3.10" dependencies = [ "datasets>=4.0.0", "fastapi>=0.117.1", - "files-to-prompt>=0.6", "psutil>=7.1.0", + "python-dotenv>=1.2.1", "regex>=2025.9.1", + "rustbpe>=0.1.0", "setuptools>=80.9.0", "tiktoken>=0.11.0", "tokenizers>=0.22.0", "torch>=2.8.0", + "transformers>=4.57.3", "uvicorn>=0.36.0", "wandb>=0.21.3", ] -[build-system] -requires = ["maturin>=1.7,<2.0"] -build-backend = "maturin" - -[tool.maturin] -module-name = "rustbpe" -bindings = "pyo3" -python-source = "." -manifest-path = "rustbpe/Cargo.toml" - [dependency-groups] dev = [ - "maturin>=1.9.4", "pytest>=8.0.0", ] @@ -45,33 +36,33 @@ python_functions = ["test_*"] # target torch to cuda 12.8 or CPU [tool.uv.sources] -torch = [ - { index = "pytorch-cpu", extra = "cpu" }, - { index = "pytorch-cu128", extra = "gpu" }, +torch = [ + { index = "pytorch-cpu", extra = "cpu" }, + { index = "pytorch-cu128", extra = "gpu" }, ] -[[tool.uv.index]] -name = "pytorch-cpu" -url = "https://download.pytorch.org/whl/cpu" -explicit = true - -[[tool.uv.index]] -name = "pytorch-cu128" -url = "https://download.pytorch.org/whl/cu128" +[[tool.uv.index]] +name = "pytorch-cpu" +url = "https://download.pytorch.org/whl/cpu" explicit = true -[project.optional-dependencies] -cpu = [ - "torch>=2.8.0", -] -gpu = [ - "torch>=2.8.0", -] - -[tool.uv] -conflicts = [ - [ - { extra = "cpu" }, - { extra = "gpu" }, - ], -] \ No newline at end of file +[[tool.uv.index]] +name = "pytorch-cu128" +url = "https://download.pytorch.org/whl/cu128" +explicit = true + +[project.optional-dependencies] +cpu = [ + "torch>=2.9.1", +] +gpu = [ + "torch>=2.9.1", +] + +[tool.uv] +conflicts = [ + [ + { extra = "cpu" }, + { extra = "gpu" }, + ], +] diff --git a/run1000.sh b/run1000.sh index 58ee3bc1..a0a66067 100644 --- a/run1000.sh +++ b/run1000.sh @@ -16,9 +16,6 @@ if [ -z "$WANDB_RUN" ]; then WANDB_RUN=dummy fi python -m nanochat.report reset -curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y -source "$HOME/.cargo/env" -uv run maturin develop --release --manifest-path rustbpe/Cargo.toml curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl # train tokenizer on ~4B characters and kick off download of the rest for pretraining diff --git a/rustbpe/Cargo.lock b/rustbpe/Cargo.lock deleted file mode 100644 index 69f8754c..00000000 --- a/rustbpe/Cargo.lock +++ /dev/null @@ -1,458 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 4 - -[[package]] -name = "ahash" -version = "0.8.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" -dependencies = [ - "cfg-if", - "getrandom", - "once_cell", - "version_check", - "zerocopy", -] - -[[package]] -name = "aho-corasick" -version = "1.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" -dependencies = [ - "memchr", -] - -[[package]] -name = "arc-swap" -version = "1.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457" - -[[package]] -name = "autocfg" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" - -[[package]] -name = "bit-set" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" -dependencies = [ - "bit-vec", -] - -[[package]] -name = "bit-vec" -version = "0.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" - -[[package]] -name = "castaway" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" -dependencies = [ - "rustversion", -] - -[[package]] -name = "cfg-if" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9" - -[[package]] -name = "compact_str" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" -dependencies = [ - "castaway", - "cfg-if", - "itoa", - "rustversion", - "ryu", - "static_assertions", -] - -[[package]] -name = "crossbeam-deque" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" -dependencies = [ - "crossbeam-epoch", - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-epoch" -version = "0.9.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" -dependencies = [ - "crossbeam-utils", -] - -[[package]] -name = "crossbeam-utils" -version = "0.8.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" - -[[package]] -name = "dary_heap" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04d2cd9c18b9f454ed67da600630b021a8a80bf33f8c95896ab33aaf1c26b728" - -[[package]] -name = "either" -version = "1.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" - -[[package]] -name = "equivalent" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" - -[[package]] -name = "fancy-regex" -version = "0.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf04c5ec15464ace8355a7b440a33aece288993475556d461154d7a62ad9947c" -dependencies = [ - "bit-set", - "regex-automata", - "regex-syntax", -] - -[[package]] -name = "getrandom" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" -dependencies = [ - "cfg-if", - "libc", - "r-efi", - "wasi", -] - -[[package]] -name = "hashbrown" -version = "0.15.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" - -[[package]] -name = "heck" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" - -[[package]] -name = "indexmap" -version = "2.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2481980430f9f78649238835720ddccc57e52df14ffce1c6f37391d61b563e9" -dependencies = [ - "equivalent", - "hashbrown", -] - -[[package]] -name = "indoc" -version = "2.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd" - -[[package]] -name = "itoa" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" - -[[package]] -name = "libc" -version = "0.2.175" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543" - -[[package]] -name = "log" -version = "0.4.28" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432" - -[[package]] -name = "memchr" -version = "2.7.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" - -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - -[[package]] -name = "once_cell" -version = "1.21.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" - -[[package]] -name = "portable-atomic" -version = "1.11.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483" - -[[package]] -name = "proc-macro2" -version = "1.0.101" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" -dependencies = [ - "unicode-ident", -] - -[[package]] -name = "pyo3" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872" -dependencies = [ - "cfg-if", - "indoc", - "libc", - "memoffset", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", - "unindent", -] - -[[package]] -name = "pyo3-build-config" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb" -dependencies = [ - "once_cell", - "target-lexicon", -] - -[[package]] -name = "pyo3-ffi" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d" -dependencies = [ - "libc", - "pyo3-build-config", -] - -[[package]] -name = "pyo3-log" -version = "0.12.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45192e5e4a4d2505587e27806c7b710c231c40c56f3bfc19535d0bb25df52264" -dependencies = [ - "arc-swap", - "log", - "pyo3", -] - -[[package]] -name = "pyo3-macros" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da" -dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn", -] - -[[package]] -name = "pyo3-macros-backend" -version = "0.23.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028" -dependencies = [ - "heck", - "proc-macro2", - "pyo3-build-config", - "quote", - "syn", -] - -[[package]] -name = "quote" -version = "1.0.40" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "r-efi" -version = "5.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" - -[[package]] -name = "rayon" -version = "1.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" -dependencies = [ - "either", - "rayon-core", -] - -[[package]] -name = "rayon-core" -version = "1.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" -dependencies = [ - "crossbeam-deque", - "crossbeam-utils", -] - -[[package]] -name = "regex-automata" -version = "0.4.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b9458fa0bfeeac22b5ca447c63aaf45f28439a709ccd244698632f9aa6394d6" -dependencies = [ - "aho-corasick", - "memchr", - "regex-syntax", -] - -[[package]] -name = "regex-syntax" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" - -[[package]] -name = "rustbpe" -version = "0.1.0" -dependencies = [ - "ahash", - "compact_str", - "dary_heap", - "fancy-regex", - "indexmap", - "log", - "pyo3", - "pyo3-log", - "rayon", -] - -[[package]] -name = "rustversion" -version = "1.0.22" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" - -[[package]] -name = "ryu" -version = "1.0.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" - -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - -[[package]] -name = "syn" -version = "2.0.106" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" -dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", -] - -[[package]] -name = "target-lexicon" -version = "0.12.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" - -[[package]] -name = "unicode-ident" -version = "1.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" - -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - -[[package]] -name = "version_check" -version = "0.9.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" - -[[package]] -name = "wasi" -version = "0.14.4+wasi-0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88a5f4a424faf49c3c2c344f166f0662341d470ea185e939657aaff130f0ec4a" -dependencies = [ - "wit-bindgen", -] - -[[package]] -name = "wit-bindgen" -version = "0.45.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c573471f125075647d03df72e026074b7203790d41351cd6edc96f46bcccd36" - -[[package]] -name = "zerocopy" -version = "0.8.26" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f" -dependencies = [ - "zerocopy-derive", -] - -[[package]] -name = "zerocopy-derive" -version = "0.8.26" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] diff --git a/rustbpe/Cargo.toml b/rustbpe/Cargo.toml deleted file mode 100644 index 392a8283..00000000 --- a/rustbpe/Cargo.toml +++ /dev/null @@ -1,15 +0,0 @@ -[package] -name = "rustbpe" -version = "0.1.0" -edition = "2024" - -[dependencies] -dary_heap = "0.3" -indexmap = "2.2" -fancy-regex = "0.16.1" -log = "0.4.28" -pyo3 = { version = "0.23.3", features = ["extension-module"] } -pyo3-log = "0.12.4" -ahash = "0.8.12" -rayon = "1.11.0" -compact_str = "0.9.0" diff --git a/rustbpe/README.md b/rustbpe/README.md deleted file mode 100644 index c88636c7..00000000 --- a/rustbpe/README.md +++ /dev/null @@ -1,5 +0,0 @@ -# rustbpe - -> The missing tiktoken training code - -A very lightweight Rust library for training a GPT tokenizer. The issue is that the inference library [tiktoken](https://github.com/openai/tiktoken) is great, but only does inference. Separately, the huggingface [tokenizers](https://github.com/huggingface/tokenizers) library does training, but it is rather bloated and really hard to navigate because it has to support all the different historical baggage of how people dealt with tokenizers over the years. More recently, I also wrote the [minbpe](https://github.com/karpathy/minbpe) library which does both training and inference, but only in inefficient Python. Basically what I really want is a non-fancy, super simple, but still relatively efficient training code for GPT tokenizer (more efficient than minbpe, much cleaner/simpler than tokenizers), and then export the trained vocab for inference with tiktoken. Does that make sense? So here we are. There are more opportunities for optimization here, I just stopped a bit early because unlike minbpe before it, rustbpe is now simple and fast enough, and not a significant bottleneck for nanochat. diff --git a/rustbpe/src/lib.rs b/rustbpe/src/lib.rs deleted file mode 100644 index 273d7f25..00000000 --- a/rustbpe/src/lib.rs +++ /dev/null @@ -1,475 +0,0 @@ -use std::cmp::Ordering; -use std::collections::HashMap as StdHashMap; - -use dary_heap::OctonaryHeap; -use fancy_regex::Regex; -use pyo3::prelude::*; - -use ahash::{AHashMap, AHashSet}; -use compact_str::CompactString; -use rayon::prelude::*; - -// Default GPT-4 style regex pattern for splitting text -const GPT4_PATTERN: &str = r"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"; - -type Pair = (u32, u32); - -/// A Byte Pair Encoding tokenizer that matches the GPT-4 style implementation -#[pyclass] -pub struct Tokenizer { - /// Maps pairs of token IDs to their merged token ID - pub merges: StdHashMap, - /// The regex pattern used for text splitting - pub pattern: String, - /// Compiled regex for efficiency - compiled_pattern: Regex, -} - -// ------------------------ internal helpers ------------------------ - -#[derive(Clone, Debug)] -struct Word { - ids: Vec, -} - -impl Word { - #[inline] - fn new(ids: Vec) -> Self { - Self { ids } - } - - #[inline] - fn pairs<'a>(&'a self) -> impl Iterator + 'a { - self.ids.windows(2).map(|w| (w[0], w[1])) - } - - /// Merge all non-overlapping occurrences of pair -> new_id. - /// Returns a small Vec of local pair-count deltas for THIS word only: - /// -1 for removed pairs, +1 for newly created pairs. - /// - /// NOTE: this version deliberately avoids a HashMap in the hot loop. - fn merge_pair(&mut self, pair: Pair, new_id: u32) -> Vec<(Pair, i32)> { - let (a, b) = pair; - let n = self.ids.len(); - if n < 2 { - return Vec::new(); - } - - let mut out: Vec = Vec::with_capacity(n); - let mut deltas: Vec<(Pair, i32)> = Vec::with_capacity(6); - - let mut i = 0; - while i < n { - if i + 1 < n && self.ids[i] == a && self.ids[i + 1] == b { - let left = out.last().copied(); - let right = if i + 2 < n { Some(self.ids[i + 2]) } else { None }; - - // remove old pairs - if let Some(x) = left { - deltas.push(((x, a), -1)); - deltas.push(((x, new_id), 1)); - } - deltas.push(((a, b), -1)); - if let Some(y) = right { - deltas.push(((b, y), -1)); - deltas.push(((new_id, y), 1)); - } - - // write merged token - out.push(new_id); - i += 2; // skip 'a' and 'b' - } else { - out.push(self.ids[i]); - i += 1; - } - } - - self.ids = out; - deltas - } -} - -#[derive(Debug, Eq)] -struct MergeJob { - pair: Pair, - count: u64, - /// set of word indices where this pair may occur and needs processing - pos: AHashSet, -} - -impl PartialEq for MergeJob { - fn eq(&self, other: &Self) -> bool { - self.count == other.count && self.pair == other.pair - } -} - -impl PartialOrd for MergeJob { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for MergeJob { - fn cmp(&self, other: &Self) -> Ordering { - // Max-heap by count; tie-break to ascending pair order (deterministic) - if self.count != other.count { - self.count.cmp(&other.count) - } else { - // ascending order on the pair when counts tie - other.pair.cmp(&self.pair) - } - } -} - -#[inline] -fn count_pairs_parallel( - words: &[Word], - counts: &[i32], -) -> (AHashMap, AHashMap>) { - words - .par_iter() - .enumerate() - .map(|(i, w)| { - let mut local_pc: AHashMap = AHashMap::new(); - let mut local_wtu: AHashMap> = AHashMap::new(); - if w.ids.len() >= 2 && counts[i] != 0 { - for (a, b) in w.pairs() { - *local_pc.entry((a, b)).or_default() += counts[i]; - local_wtu.entry((a, b)).or_default().insert(i); - } - } - (local_pc, local_wtu) - }) - .reduce( - || (AHashMap::new(), AHashMap::new()), - |(mut acc_pc, mut acc_wtu), (pc, wtu)| { - for (k, v) in pc { - *acc_pc.entry(k).or_default() += v; - } - for (k, s) in wtu { - acc_wtu.entry(k).or_default().extend(s); - } - (acc_pc, acc_wtu) - }, - ) -} - -// ------------------------ END helpers ------------------------ - -impl Tokenizer { - - /// Core incremental BPE training given unique words and their counts. - /// `words`: one entry per unique chunk (Vec of token-ids/bytes). - /// `counts`: same length as `words`, count per chunk. - fn train_core_incremental(&mut self, mut words: Vec, counts: Vec, vocab_size: u32) { - assert!(vocab_size >= 256, "vocab_size must be at least 256"); - let num_merges = vocab_size - 256; - log::info!("Starting BPE training: {} merges to compute", num_merges); - self.merges.clear(); - - // ---- Initial pair_counts and where_to_update (parallel) ---- - log::info!("Computing initial pair counts from {} unique sequences", words.len()); - let (mut pair_counts, mut where_to_update) = count_pairs_parallel(&words, &counts); - - // ---- Build heap ---- - log::info!("Building heap with {} unique pairs", pair_counts.len()); - let mut heap = OctonaryHeap::with_capacity(pair_counts.len()); - for (pair, pos) in where_to_update.drain() { - let c = *pair_counts.get(&pair).unwrap_or(&0); - if c > 0 { - heap.push(MergeJob { - pair, - count: c as u64, - pos, - }); - } - } - - // ---- Merge loop ---- - log::info!("Starting merge loop"); - let mut merges_done = 0u32; - let mut last_log_percent = 0u32; - - while merges_done < num_merges { - let Some(mut top) = heap.pop() else { break; }; - - // Lazy refresh - let current = *pair_counts.get(&top.pair).unwrap_or(&0); - if top.count != current as u64 { - top.count = current as u64; - if top.count > 0 { - heap.push(top); - } - continue; - } - if top.count == 0 { - break; - } - - // Record merge - let new_id = 256 + merges_done; - self.merges.insert(top.pair, new_id); - - // Merge this pair in all words where it occurs - let mut local_pos_updates: AHashMap> = AHashMap::new(); - for &word_idx in &top.pos { - // Apply merge to this word and collect pair-count deltas - let changes = words[word_idx].merge_pair(top.pair, new_id); - // Update global pair counts based on this word's count - for (pair, delta) in changes { - let delta_total = delta * counts[word_idx]; - if delta_total != 0 { - *pair_counts.entry(pair).or_default() += delta_total; - if delta > 0 { - local_pos_updates.entry(pair).or_default().insert(word_idx); - } - } - } - } - - // Add the updated pair counts back to the heap - for (pair, pos) in local_pos_updates { - let cnt = *pair_counts.get(&pair).unwrap_or(&0); - if cnt > 0 { - heap.push(MergeJob { - pair, - count: cnt as u64, - pos, - }); - } - } - - merges_done += 1; - - // Log progress every 1% - let current_percent = (merges_done * 100) / num_merges; - if current_percent > last_log_percent { - log::info!( - "Progress: {}% ({}/{} merges) - Last merge: {:?} -> {} (frequency: {})", - current_percent, merges_done, num_merges, top.pair, new_id, top.count - ); - last_log_percent = current_percent; - } - } - - log::info!("Finished training: {} merges completed", merges_done); - } -} - -/// Public methods for the Tokenizer class that will be exposed to Python. -#[pymethods] -impl Tokenizer { - /// Create a new Tokenizer - #[new] - pub fn new() -> Self { - Self { - merges: StdHashMap::new(), - pattern: String::new(), - compiled_pattern: Regex::new("").expect("Empty regex should be valid"), - } - } - - /// Train from a streaming iterator (parallel ingestion). - /// We refill a Rust Vec buffer under the GIL, then release the GIL - /// to do the heavy splitting and counting **in parallel** with rayon. - #[pyo3(signature = (iterator, vocab_size, buffer_size=8192, pattern=None))] - #[pyo3(text_signature = "(self, iterator, vocab_size, buffer_size=8192, pattern=None)")] - pub fn train_from_iterator( - &mut self, - py: pyo3::Python<'_>, - iterator: &pyo3::Bound<'_, pyo3::PyAny>, - vocab_size: u32, - buffer_size: usize, - pattern: Option, - ) -> PyResult<()> { - // Use provided pattern or default to GPT-4 pattern - let pattern_str = pattern.unwrap_or_else(|| GPT4_PATTERN.to_string()); - - // Update the stored pattern and compile it - self.pattern = pattern_str.clone(); - self.compiled_pattern = Regex::new(&pattern_str) - .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid regex pattern: {}", e)))?; - - // Prepare a true Python iterator object - let py_iter: pyo3::Py = unsafe { - pyo3::Py::from_owned_ptr_or_err(py, pyo3::ffi::PyObject_GetIter(iterator.as_ptr()))? - }; - - // Global chunk counts - let mut counts: AHashMap = AHashMap::new(); - - // Temporary buffer we refill under the GIL - let mut buf: Vec = Vec::with_capacity(buffer_size); - - log::info!("Processing sequences from iterator (buffer_size: {})", buffer_size); - let mut total_sequences = 0u64; - - // Helper: refill `buf` with up to `buffer_size` strings from the Python iterator. - // Returns Ok(true) if the iterator is exhausted, Ok(false) otherwise. - let refill = |buf: &mut Vec| -> PyResult { - pyo3::Python::with_gil(|py| { - buf.clear(); - let it = py_iter.bind(py); - loop { - if buf.len() >= buffer_size { - return Ok(false); - } - // next(it) - let next_obj = unsafe { - pyo3::Bound::from_owned_ptr_or_opt(py, pyo3::ffi::PyIter_Next(it.as_ptr())) - }; - match next_obj { - Some(obj) => { - let s: String = obj.extract()?; - buf.push(s); - } - None => { - if pyo3::PyErr::occurred(py) { - return Err(pyo3::PyErr::fetch(py)); - } else { - return Ok(true); // exhausted - } - } - } - } - }) - }; - - // Stream ingestion loop: refill under GIL, process without GIL (parallel) - loop { - let exhausted = refill(&mut buf)?; - if buf.is_empty() && exhausted { - break; - } - - total_sequences += buf.len() as u64; - - let pattern = self.compiled_pattern.clone(); - let local: AHashMap = py.allow_threads(|| { - buf.par_iter() - .map(|s| { - let mut m: AHashMap = AHashMap::new(); - for mat in pattern.find_iter(s) { - let piece = mat.expect("regex match failed").as_str(); - *m.entry(CompactString::from(piece)).or_default() += 1; - } - m - }) - .reduce( - || AHashMap::new(), - |mut a, b| { - for (k, v) in b { - *a.entry(k).or_default() += v; - } - a - }, - ) - }); - - // Merge local into global (single-threaded) - for (k, v) in local { - *counts.entry(k).or_default() += v; - } - - if exhausted { - break; - } - } - log::info!("Processed {} sequences total, {} unique", total_sequences, counts.len()); - - // Materialize words & counts - let mut words = Vec::with_capacity(counts.len()); - let mut cvec = Vec::with_capacity(counts.len()); - for (chunk, c) in counts.into_iter() { - words.push(Word::new(chunk.as_bytes().iter().map(|&b| b as u32).collect())); - cvec.push(c); - } - - self.train_core_incremental(words, cvec, vocab_size); - Ok(()) - } - - /// Return the regex pattern - pub fn get_pattern(&self) -> String { - self.pattern.clone() - } - - /// Return the mergeable ranks (token bytes -> token id / rank) - pub fn get_mergeable_ranks(&self) -> Vec<(Vec, u32)> { - let mut mergeable_ranks = Vec::new(); - - // Build vocabulary incrementally from low to high token IDs - let mut token_bytes: Vec> = (0..256_u32).map(|i| vec![i as u8]).collect(); - - for (i, bytes) in token_bytes.iter().enumerate() { - mergeable_ranks.push((bytes.clone(), i as u32)); - } - - // Sort merges by token id (so we can reconstruct bytes progressively) - let mut sorted_merges: Vec<_> = self.merges.iter().collect(); - sorted_merges.sort_by_key(|&(_, &token_id)| token_id); - - for (&pair, &merged_id) in sorted_merges { - let (left, right) = pair; - let mut merged_bytes = token_bytes[left as usize].clone(); - merged_bytes.extend(&token_bytes[right as usize]); - - if token_bytes.len() <= merged_id as usize { - token_bytes.resize(merged_id as usize + 1, Vec::new()); - } - token_bytes[merged_id as usize] = merged_bytes.clone(); - - mergeable_ranks.push((merged_bytes, merged_id)); - } - - mergeable_ranks - } - - /// Encode a string into token IDs - pub fn encode(&self, text: &str) -> Vec { - let mut all_ids = Vec::new(); - - // Split text using the regex pattern - for m in self.compiled_pattern.find_iter(text) { - let chunk = m.expect("regex match failed").as_str(); - - // Convert chunk to bytes then to u32 IDs - let mut ids: Vec = chunk.bytes().map(|b| b as u32).collect(); - - // Apply merges iteratively - while ids.len() >= 2 { - // Find the best pair to merge - let mut best_pair: Option<(usize, Pair, u32)> = None; - - for i in 0..ids.len() - 1 { - let pair: Pair = (ids[i], ids[i + 1]); - if let Some(&new_id) = self.merges.get(&pair) { - if best_pair.is_none() || new_id < best_pair.unwrap().2 { - best_pair = Some((i, pair, new_id)); - } - } - } - - // If we found a pair to merge, apply it - if let Some((idx, _pair, new_id)) = best_pair { - ids[idx] = new_id; - ids.remove(idx + 1); - } else { - // No more merges possible - break; - } - } - - all_ids.extend(ids); - } - - all_ids - } -} - -#[pymodule] -fn rustbpe(m: &Bound<'_, PyModule>) -> PyResult<()> { - pyo3_log::init(); // forwards Rust `log` to Python's `logging` - m.add_class::()?; - Ok(()) -} diff --git a/scripts/base_eval.py b/scripts/base_eval.py index 36635382..bd83ff36 100644 --- a/scripts/base_eval.py +++ b/scripts/base_eval.py @@ -149,6 +149,8 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument('--hf-path', type=str, default=None, help='HuggingFace model path to evaluate') parser.add_argument('--max-per-task', type=int, default=-1, help='Max examples per task to evaluate (-1 = disable)') + parser.add_argument('--model-tag', type=str, default=None, help='optional model tag for the output directory name') + parser.add_argument('--step', type=str, default=None, help='optional model step for the output directory name') args = parser.parse_args() # distributed / precision setup @@ -166,7 +168,7 @@ def main(): model_slug = hf_path.replace("/", "-") # for the output csv file else: # load a local model from the file system - model, tokenizer, meta = load_model("base", device, phase="eval") + model, tokenizer, meta = load_model("base", device, phase="eval", model_tag=args.model_tag, step=args.step) model_name = f"base_model (step {meta['step']})" # just for logging model_slug = f"base_model_{meta['step']:06d}" # for the output csv file diff --git a/scripts/base_loss.py b/scripts/base_loss.py index abcde5f9..3dbe68f8 100644 --- a/scripts/base_loss.py +++ b/scripts/base_loss.py @@ -6,7 +6,7 @@ Loads a checkpoint, and: Example run as: torchrun --standalone --nproc_per_node=8 -m scripts.base_loss """ -import os +import argparse from contextlib import nullcontext import torch from nanochat.checkpoint_manager import load_model @@ -16,29 +16,30 @@ from nanochat.tokenizer import get_token_bytes from nanochat.loss_eval import evaluate_bpb from nanochat.engine import Engine -# Configuration -device_batch_size = 32 -split_tokens = 20*524288 # number of tokens to evaluate per split -model_tag = None # optional model tag for the output directory name -model_step = None # optional model step for the output directory name -device_type = "" # cuda|cpu|mps (empty => autodetect) -exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file +# CLI arguments +parser = argparse.ArgumentParser(description="Evaluate loss on train/val splits and sample from model") +parser.add_argument("--device_batch_size", type=int, default=32, help="per-device batch size") +parser.add_argument("--split_tokens", type=int, default=20*524288, help="number of tokens to evaluate per split") +parser.add_argument("--model_tag", type=str, default=None, help="model tag for checkpoint directory") +parser.add_argument("--model_step", type=int, default=None, help="model step to load") +parser.add_argument("--device_type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") +args = parser.parse_args() # Load the base model and the tokenizer -device_type = autodetect_device_type() if device_type == "" else device_type +device_type = autodetect_device_type() if args.device_type == "" else args.device_type ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) -model, tokenizer, meta = load_model("base", device, phase="eval", model_tag=model_tag, step=model_step) +model, tokenizer, meta = load_model("base", device, phase="eval", model_tag=args.model_tag, step=args.model_step) sequence_len = meta["model_config"]["sequence_len"] # could be arbitrary really autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext() # Evaluate the loss on each split -tokens_per_step = device_batch_size * sequence_len * ddp_world_size -assert split_tokens % tokens_per_step == 0, "split_tokens must be divisible by tokens_per_step" -steps = split_tokens // tokens_per_step +tokens_per_step = args.device_batch_size * sequence_len * ddp_world_size +assert args.split_tokens % tokens_per_step == 0, "split_tokens must be divisible by tokens_per_step" +steps = args.split_tokens // tokens_per_step token_bytes = get_token_bytes(device=device) bpb_results = {} for split_name in ["train", "val"]: - loader = tokenizing_distributed_data_loader(device_batch_size, sequence_len, split_name, device=device) + loader = tokenizing_distributed_data_loader(args.device_batch_size, sequence_len, split_name, device=device) with autocast_ctx: bpb = evaluate_bpb(model, loader, steps, token_bytes) print0(f"{split_name} bpb: {bpb:.4f}") diff --git a/scripts/base_train.py b/scripts/base_train.py index c9ea6c99..2390b688 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -12,7 +12,8 @@ python -m scripts.base_train --depth=4 --max_seq_len=512 --device_batch_size=1 - """ import os -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" +os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" +import argparse import time from contextlib import nullcontext @@ -30,46 +31,46 @@ from scripts.base_eval import evaluate_model print_banner() # ----------------------------------------------------------------------------- -# User settings -run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb) +# CLI arguments +parser = argparse.ArgumentParser(description="Pretrain base model") +# Logging +parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") # Runtime -device_type = "" # cuda|cpu|mps (empty => autodetect good device type default, in order: CUDA > MPS > CPU) +parser.add_argument("--device_type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") # Model architecture -depth = 20 # the depth of the Transformer model to train, rest of the kwargs are derived -max_seq_len = 2048 # max context length -# Training horizon. Only one of these 3 will be used, in this order of precedence. -num_iterations = -1 # explicit number of steps of the optimization (-1 = disable) -target_flops = -1.0 # calculate num_iterations to reach target_flops. Useful for scaling laws experiments (-1 = disable) -target_param_data_ratio = 20 # calculate num_iterations to maintain fixed data:param ratio (Chinchilla=20) (-1 = disable) +parser.add_argument("--depth", type=int, default=20, help="depth of the Transformer model") +parser.add_argument("--max_seq_len", type=int, default=2048, help="max context length") +# Training horizon (only one used, in order of precedence) +parser.add_argument("--num_iterations", type=int, default=-1, help="explicit number of optimization steps (-1 = disable)") +parser.add_argument("--target_flops", type=float, default=-1.0, help="calculate num_iterations to reach target_flops (-1 = disable)") +parser.add_argument("--target_param_data_ratio", type=int, default=20, help="calculate num_iterations to maintain data:param ratio (Chinchilla=20, -1 = disable)") # Optimization -device_batch_size = 32 # per-device batch size (set to not OOM) -total_batch_size = 524288 # total desired batch size, in #tokens -embedding_lr = 0.2 # learning rate for the embedding parameters (Adam) -unembedding_lr = 0.004 # learning rate for the unembedding parameters (Adam) -weight_decay = 0.0 # weight decay for the embedding/unembedding parameters (Adam) -matrix_lr = 0.02 # learning rate for the matrix parameters (Muon) -grad_clip = 1.0 # gradient clipping value (0.0 = disabled) -warmup_ratio = 0.0 # ratio of iterations for LR warmup -warmdown_ratio = 0.2 # ratio of iterations for LR warmdown -final_lr_frac = 0.0 # final LR is this fraction of the initial LR -resume_from_step = -1 # resume training from this step of the optimization (-1 = disable) +parser.add_argument("--device_batch_size", type=int, default=32, help="per-device batch size") +parser.add_argument("--total_batch_size", type=int, default=524288, help="total batch size in tokens") +parser.add_argument("--embedding_lr", type=float, default=0.2, help="learning rate for embedding parameters (Adam)") +parser.add_argument("--unembedding_lr", type=float, default=0.004, help="learning rate for unembedding parameters (Adam)") +parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay for embedding/unembedding parameters (Adam)") +parser.add_argument("--matrix_lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)") +parser.add_argument("--grad_clip", type=float, default=1.0, help="gradient clipping value (0.0 = disabled)") +parser.add_argument("--warmup_ratio", type=float, default=0.0, help="ratio of iterations for LR warmup") +parser.add_argument("--warmdown_ratio", type=float, default=0.2, help="ratio of iterations for LR warmdown") +parser.add_argument("--final_lr_frac", type=float, default=0.0, help="final LR as fraction of initial LR") +parser.add_argument("--resume_from_step", type=int, default=-1, help="resume training from this step (-1 = disable)") # Evaluation -eval_every = 250 # every how many steps to evaluate the model for val bpb -eval_tokens = 20*524288 # number of tokens to evaluate val loss on -core_metric_every = 2000 # every how many steps to evaluate the core metric (-1 = disable) -core_metric_max_per_task = 500 # examples per task in estimating the core metric -sample_every = 2000 # every how many steps to sample from the model -save_every = -1 # every how many steps to save model checkpoints (-1 = disable, and save only at the end of the run) +parser.add_argument("--eval_every", type=int, default=250, help="evaluate val bpb every N steps (-1 = disable)") +parser.add_argument("--eval_tokens", type=int, default=20*524288, help="number of tokens to evaluate val loss on") +parser.add_argument("--core_metric_every", type=int, default=2000, help="evaluate CORE metric every N steps (-1 = disable)") +parser.add_argument("--core_metric_max_per_task", type=int, default=500, help="examples per task for CORE metric") +parser.add_argument("--sample_every", type=int, default=2000, help="sample from model every N steps (-1 = disable)") +parser.add_argument("--save_every", type=int, default=-1, help="save checkpoints every N steps (-1 = only at end)") # Output -model_tag = "" # optionally override the model tag for the output checkpoint directory name -# now allow CLI to override the settings via the configurator lol -config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))] -exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file -user_config = {k: globals()[k] for k in config_keys} # will be useful for logging +parser.add_argument("--model_tag", type=str, default=None, help="override model tag for checkpoint directory name") +args = parser.parse_args() +user_config = vars(args).copy() # for logging # ----------------------------------------------------------------------------- # Compute init -device_type = autodetect_device_type() if device_type == "" else device_type +device_type = autodetect_device_type() if args.device_type == "" else args.device_type ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext() @@ -77,8 +78,8 @@ synchronize = torch.cuda.synchronize if device_type == "cuda" else lambda: None get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0 # wandb logging init -use_dummy_wandb = run == "dummy" or not master_process -wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat", name=run, config=user_config) +use_dummy_wandb = args.run == "dummy" or not master_process +wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat", name=args.run, config=user_config) # Tokenizer will be useful for evaluation, also we need the vocab size tokenizer = get_tokenizer() @@ -87,9 +88,17 @@ vocab_size = tokenizer.get_vocab_size() print0(f"Vocab size: {vocab_size:,}") # Model kwargs are derived from the desired depth of the model -num_layers = depth -model_dim = depth * 64 # aspect ratio 64 (usually this is varied from 64 -> 128 as model size increases) -num_heads = max(1, (model_dim + 127) // 128) # head dim 128 (the division here is ceil div) +num_layers = args.depth +model_dim = args.depth * 64 # aspect ratio 64 (usually this is varied from 64 -> 128 as model size increases) +def find_num_heads(model_dim, target_head_dim=128): + # Find num_heads that divides model_dim evenly, with head_dim closest to target. + ideal = max(1, round(model_dim / target_head_dim)) + for offset in range(model_dim): + for candidate in [ideal + offset, ideal - offset]: + if candidate > 0 and model_dim % candidate == 0: + return candidate + return 1 +num_heads = find_num_heads(model_dim) num_kv_heads = num_heads # default is 1:1 GQA (Group Query Attention) ratio (i.e. GQA is disabled) print0(f"num_layers: {num_layers}") print0(f"model_dim: {model_dim}") @@ -98,33 +107,34 @@ print0(f"num_kv_heads: {num_kv_heads}") # Optimizer / data / training length related hyperparameters # figure out the needed gradient accumulation to reach the desired total batch size -tokens_per_fwdbwd = device_batch_size * max_seq_len # tokens per iteration for a single rank +tokens_per_fwdbwd = args.device_batch_size * args.max_seq_len # tokens per iteration for a single rank world_tokens_per_fwdbwd = tokens_per_fwdbwd * ddp_world_size # total tokens per iteration for all ranks -assert total_batch_size % world_tokens_per_fwdbwd == 0 -grad_accum_steps = total_batch_size // world_tokens_per_fwdbwd -print0(f"Tokens / micro-batch / rank: {device_batch_size} x {max_seq_len} = {tokens_per_fwdbwd:,}") +assert args.total_batch_size % world_tokens_per_fwdbwd == 0 +grad_accum_steps = args.total_batch_size // world_tokens_per_fwdbwd +print0(f"Tokens / micro-batch / rank: {args.device_batch_size} x {args.max_seq_len} = {tokens_per_fwdbwd:,}") print0(f"Tokens / micro-batch: {world_tokens_per_fwdbwd:,}") -print0(f"Total batch size {total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}") +print0(f"Total batch size {args.total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}") # ----------------------------------------------------------------------------- # Initialize the Model # Create a new model with random weights -model_config_kwargs = dict(sequence_len=max_seq_len, vocab_size=vocab_size, n_layer=num_layers, n_head=num_heads, n_kv_head=num_kv_heads, n_embd=model_dim) +model_config_kwargs = dict(sequence_len=args.max_seq_len, vocab_size=vocab_size, n_layer=num_layers, n_head=num_heads, n_kv_head=num_kv_heads, n_embd=model_dim) with torch.device("meta"): + # All tensors are created as meta tensors (they have shape/dtype but no data) model_config = GPTConfig(**model_config_kwargs) model = GPT(model_config) -model.to_empty(device=device) -model.init_weights() +model.to_empty(device=device) # All tensors get storage on target device but with uninitialized (garbage) data +model.init_weights() # All tensors get initialized # If we are resuming, overwrite the model parameters with those of the checkpoint base_dir = get_base_dir() -output_dirname = model_tag if model_tag else f"d{depth}" # e.g. d12 +output_dirname = args.model_tag if args.model_tag else f"d{args.depth}" # e.g. d12 checkpoint_dir = os.path.join(base_dir, "base_checkpoints", output_dirname) -resuming = resume_from_step != -1 +resuming = args.resume_from_step != -1 if resuming: - print0(f"Resuming optimization from step {resume_from_step}") - model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, resume_from_step, device, load_optimizer=True, rank=ddp_rank) + print0(f"Resuming optimization from step {args.resume_from_step}") + model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, args.resume_from_step, device, load_optimizer=True, rank=ddp_rank) model.load_state_dict(model_data, strict=True, assign=True) del model_data # free up this memory after the copy @@ -136,28 +146,29 @@ num_flops_per_token = model.estimate_flops() print0(f"Estimated FLOPs per token: {num_flops_per_token:e}") # Calculate number of iterations. Either it is given, or from target flops, or from target data:param ratio (in that order) -assert num_iterations > 0 or target_param_data_ratio > 0 or target_flops > 0 -if num_iterations > 0: +assert args.num_iterations > 0 or args.target_param_data_ratio > 0 or args.target_flops > 0 +if args.num_iterations > 0: + num_iterations = args.num_iterations print0(f"Using user-provided number of iterations: {num_iterations:,}") -elif target_flops > 0: +elif args.target_flops > 0: # calculate the number of iterations from the target flops - num_iterations = round(target_flops / (num_flops_per_token * total_batch_size)) + num_iterations = round(args.target_flops / (num_flops_per_token * args.total_batch_size)) print0(f"Calculated number of iterations from target FLOPs: {num_iterations:,}") -elif target_param_data_ratio > 0: +elif args.target_param_data_ratio > 0: # calculate the number of iterations from the target param data ratio - target_tokens = target_param_data_ratio * num_params - num_iterations = target_tokens // total_batch_size + target_tokens = args.target_param_data_ratio * num_params + num_iterations = target_tokens // args.total_batch_size print0(f"Calculated number of iterations from target data:param ratio: {num_iterations:,}") else: raise ValueError("No training horizon specified") -total_tokens = total_batch_size * num_iterations +total_tokens = args.total_batch_size * num_iterations print0(f"Total number of training tokens: {total_tokens:,}") -print0(f"Tokens : Params ratio: {total_batch_size * num_iterations / num_params:.2f}") # Chinchilla is ~20 +print0(f"Tokens : Params ratio: {args.total_batch_size * num_iterations / num_params:.2f}") # Chinchilla is ~20 print0(f"Total training FLOPs estimate: {num_flops_per_token * total_tokens:e}") # ----------------------------------------------------------------------------- # Initialize the Optimizer (Muon for Linear layers, AdamW for embedding and lm_head) -optimizers = model.setup_optimizers(unembedding_lr=unembedding_lr, embedding_lr=embedding_lr, matrix_lr=matrix_lr, weight_decay=weight_decay) +optimizers = model.setup_optimizers(unembedding_lr=args.unembedding_lr, embedding_lr=args.embedding_lr, matrix_lr=args.matrix_lr, weight_decay=args.weight_decay) adamw_optimizer, muon_optimizer = optimizers if resuming: @@ -169,8 +180,8 @@ if resuming: # Initialize the DataLoaders for train/val tokens_dir = os.path.join(base_dir, "tokenized_data") dataloader_resume_state_dict = None if not resuming else meta_data["dataloader_state_dict"] -train_loader = tokenizing_distributed_data_loader_with_state(device_batch_size, max_seq_len, split="train", device=device, resume_state_dict=dataloader_resume_state_dict) -build_val_loader = lambda: tokenizing_distributed_data_loader(device_batch_size, max_seq_len, split="val", device=device) +train_loader = tokenizing_distributed_data_loader_with_state(args.device_batch_size, args.max_seq_len, split="train", device=device, resume_state_dict=dataloader_resume_state_dict) +build_val_loader = lambda: tokenizing_distributed_data_loader(args.device_batch_size, args.max_seq_len, split="val", device=device) x, y, dataloader_state_dict = next(train_loader) # kick off load of the very first batch of data # ----------------------------------------------------------------------------- @@ -178,15 +189,15 @@ x, y, dataloader_state_dict = next(train_loader) # kick off load of the very fir # Learning rate scheduler def get_lr_multiplier(it): - warmup_iters = round(warmup_ratio * num_iterations) - warmdown_iters = round(warmdown_ratio * num_iterations) + warmup_iters = round(args.warmup_ratio * num_iterations) + warmdown_iters = round(args.warmdown_ratio * num_iterations) if it < warmup_iters: return (it + 1) / warmup_iters elif it <= num_iterations - warmdown_iters: return 1.0 else: progress = (num_iterations - it) / warmdown_iters - return progress * 1.0 + (1 - progress) * final_lr_frac + return progress * 1.0 + (1 - progress) * args.final_lr_frac # Momentum scheduler for Muon optimizer def get_muon_momentum(it): @@ -199,12 +210,14 @@ def get_muon_momentum(it): if not resuming: step = 0 + val_bpb = None # will be set if eval_every > 0 min_val_bpb = float("inf") smooth_train_loss = 0 # EMA of training loss total_training_time = 0 # total wall-clock time of training else: step = meta_data["step"] loop_state = meta_data["loop_state"] + val_bpb = meta_data["val_bpb"] min_val_bpb = loop_state["min_val_bpb"] smooth_train_loss = loop_state["smooth_train_loss"] total_training_time = loop_state["total_training_time"] @@ -213,13 +226,13 @@ else: # Training loop while True: last_step = step == num_iterations # loop runs num_iterations+1 times so that we can eval/save at the end - flops_so_far = num_flops_per_token * total_batch_size * step + flops_so_far = num_flops_per_token * args.total_batch_size * step # once in a while: evaluate the val bpb (all ranks participate) - if last_step or step % eval_every == 0: + if args.eval_every > 0 and (last_step or step % args.eval_every == 0): model.eval() val_loader = build_val_loader() - eval_steps = eval_tokens // (device_batch_size * max_seq_len * ddp_world_size) + eval_steps = args.eval_tokens // (args.device_batch_size * args.max_seq_len * ddp_world_size) with autocast_ctx: val_bpb = evaluate_bpb(model, val_loader, eval_steps, token_bytes) print0(f"Step {step:05d} | Validation bpb: {val_bpb:.4f}") @@ -236,10 +249,10 @@ while True: # once in a while: estimate the CORE metric (all ranks participate) # use the original uncompiled model because the inputs keep changing shape results = {} - if core_metric_every > 0 and (last_step or (step > 0 and step % core_metric_every == 0)): + if args.core_metric_every > 0 and (last_step or (step > 0 and step % args.core_metric_every == 0)): model.eval() with autocast_ctx: - results = evaluate_model(orig_model, tokenizer, device, max_per_task=core_metric_max_per_task) + results = evaluate_model(orig_model, tokenizer, device, max_per_task=args.core_metric_max_per_task) print0(f"Step {step:05d} | CORE metric: {results['core_metric']:.4f}") wandb_run.log({ "step": step, @@ -251,7 +264,7 @@ while True: # once in a while: sample from the model (only on master process) # use the original uncompiled model because the inputs keep changing shape - if master_process and (last_step or (step > 0 and step % sample_every == 0)): + if args.sample_every > 0 and master_process and (last_step or (step > 0 and step % args.sample_every == 0)): model.eval() prompts = [ "The capital of France is", @@ -271,7 +284,7 @@ while True: model.train() # save checkpoint: at the end of the run, or every save_every steps, except at the first step or the resume step - if last_step or (step > 0 and step != resume_from_step and save_every > 0 and step % save_every == 0): + if last_step or (step > 0 and step != args.resume_from_step and args.save_every > 0 and step % args.save_every == 0): save_checkpoint( checkpoint_dir, step, @@ -282,8 +295,8 @@ while True: "val_bpb": val_bpb, # loss at last step "model_config": model_config_kwargs, "user_config": user_config, # inputs to the training script - "device_batch_size": device_batch_size, - "max_seq_len": max_seq_len, + "device_batch_size": args.device_batch_size, + "max_seq_len": args.max_seq_len, "dataloader_state_dict": dataloader_state_dict, "loop_state": { # all loop state (other than step) so that we can resume training "min_val_bpb": min_val_bpb, @@ -311,9 +324,9 @@ while True: loss.backward() x, y, dataloader_state_dict = next(train_loader) # prefetch the next batch while the GPU is busy with forward/backward # gradient clipping - grad_clip_enabled = grad_clip > 0.0 + grad_clip_enabled = args.grad_clip > 0.0 if grad_clip_enabled: - grad_norm_tensor = torch.nn.utils.clip_grad_norm_(orig_model.parameters(), grad_clip) + grad_norm_tensor = torch.nn.utils.clip_grad_norm_(orig_model.parameters(), args.grad_clip) grad_norm = grad_norm_tensor.item() # GPU tensor -> CPU float (note: cpu-gpu sync point) # step the optimizers lrm = get_lr_multiplier(step) @@ -336,14 +349,23 @@ while True: smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss.item() # EMA the training loss debiased_smooth_loss = smooth_train_loss / (1 - ema_beta**(step + 1)) # debias the EMA pct_done = 100 * step / num_iterations - tok_per_sec = int(total_batch_size / dt) - flops_per_sec = num_flops_per_token * total_batch_size / dt + tok_per_sec = int(args.total_batch_size / dt) + flops_per_sec = num_flops_per_token * args.total_batch_size / dt promised_flops_per_sec_h100 = 989e12 * ddp_world_size # bfloat16 H100 SXM and without 2:4 sparsity mfu = 100 * flops_per_sec / promised_flops_per_sec_h100 # in % if step > 10: total_training_time += dt # only count the time after the first 10 steps print_grad_norm = f" grad norm: {grad_norm:.4f} |" if grad_clip_enabled else "" - print0(f"step {step:05d}/{num_iterations:05d} ({pct_done:.2f}%) | loss: {debiased_smooth_loss:.6f} |{print_grad_norm} lrm: {lrm:.2f} | dt: {dt * 1000:.2f}ms | tok/sec: {tok_per_sec:,} | mfu: {mfu:.2f} | total time: {total_training_time/60:.2f}m") + # Calculate ETA based on average time per step (excluding first 10 steps) + steps_done = step - 10 + if steps_done > 0: + avg_time_per_step = total_training_time / steps_done + remaining_steps = num_iterations - step + eta_seconds = remaining_steps * avg_time_per_step + eta_str = f" | eta: {eta_seconds/60:.1f}m" + else: + eta_str = "" + print0(f"step {step:05d}/{num_iterations:05d} ({pct_done:.2f}%) | loss: {debiased_smooth_loss:.6f} |{print_grad_norm} lrm: {lrm:.2f} | dt: {dt * 1000:.2f}ms | tok/sec: {tok_per_sec:,} | mfu: {mfu:.2f} | total time: {total_training_time/60:.2f}m{eta_str}") if step % 100 == 0: log_data = { "step": step, @@ -365,7 +387,8 @@ while True: # print a few more stats print0(f"Peak memory usage: {get_max_memory() / 1024 / 1024:.2f}MiB") print0(f"Total training time: {total_training_time/60:.2f}m") -print0(f"Minimum validation bpb: {min_val_bpb:.4f}") +if val_bpb is not None: + print0(f"Minimum validation bpb: {min_val_bpb:.4f}") # Log to report from nanochat.report import get_report @@ -376,14 +399,14 @@ get_report().log(section="Base model training", data=[ "Number of FLOPs per token": f"{num_flops_per_token:e}", "Calculated number of iterations": num_iterations, "Number of training tokens": total_tokens, - "Tokens : Params ratio": total_batch_size * num_iterations / num_params, + "Tokens : Params ratio": args.total_batch_size * num_iterations / num_params, "DDP world size": ddp_world_size, - "warmup_ratio": warmup_ratio, - "warmdown_ratio": warmdown_ratio, - "final_lr_frac": final_lr_frac, + "warmup_ratio": args.warmup_ratio, + "warmdown_ratio": args.warmdown_ratio, + "final_lr_frac": args.final_lr_frac, }, { # stats about training outcomes - "Minimum validation bpb": min_val_bpb, + "Minimum validation bpb": min_val_bpb if val_bpb is not None else None, "Final validation bpb": val_bpb, "CORE metric estimate": results.get("core_metric", None), "MFU %": f"{mfu:.2f}%", diff --git a/scripts/chat_rl.py b/scripts/chat_rl.py index bc78e797..1a09962b 100644 --- a/scripts/chat_rl.py +++ b/scripts/chat_rl.py @@ -16,55 +16,69 @@ python -m scripts.chat_rl torchrun --standalone --nproc_per_node=8 -m scripts.chat_rl -- --run=default """ +import argparse import os import itertools import re import wandb import torch import torch.distributed as dist +from contextlib import nullcontext -from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, DummyWandb +from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, DummyWandb, autodetect_device_type from nanochat.checkpoint_manager import save_checkpoint, load_model from nanochat.engine import Engine from tasks.gsm8k import GSM8K -# RL hyperparameters -run = "dummy" # wandb run name -source = "sft" # mid|sft -dtype = "bfloat16" -device_batch_size = 8 # no forward pass will go above this to not OOM -examples_per_step = 16 # in total and across all ranks (note: examples, not samples/completions!) -num_samples = 16 # number of samples per example (/question) -max_new_tokens = 256 -temperature = 1.0 -top_k = 50 # TODO: try None? -unembedding_lr = 0.004 -embedding_lr = 0.2 -matrix_lr = 0.02 -weight_decay = 0.0 -init_lr_frac = 0.05 -num_epochs = 1 # how many epochs of gsm8k to train on -save_every = 60 # every how many steps to save the model -eval_every = 60 # every how many steps to evaluate the model for val pass@k -eval_examples = 400 # number of examples used for evaluating pass@k -# now allow CLI to override the settings via the configurator lol -config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))] -exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file -user_config = {k: globals()[k] for k in config_keys} # will be useful for logging +# ----------------------------------------------------------------------------- +# CLI arguments +parser = argparse.ArgumentParser(description="Reinforcement learning on GSM8K") +# Logging +parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") +# Runtime +parser.add_argument("--device_type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") +parser.add_argument("--dtype", type=str, default="bfloat16", help="float32|bfloat16") +# Model loading +parser.add_argument("--source", type=str, default="sft", help="mid|sft - which checkpoint to load from") +parser.add_argument("--model_tag", type=str, default=None, help="model tag to load from") +parser.add_argument("--model_step", type=int, default=None, help="model step to load from") +# Training horizon +parser.add_argument("--num_epochs", type=int, default=1, help="number of epochs over GSM8K") +# Batch sizes / sampling +parser.add_argument("--device_batch_size", type=int, default=8, help="max batch size per forward pass") +parser.add_argument("--examples_per_step", type=int, default=16, help="total examples per optimization step across all ranks") +parser.add_argument("--num_samples", type=int, default=16, help="number of samples per example/question") +# Generation +parser.add_argument("--max_new_tokens", type=int, default=256, help="max tokens to generate per sample") +parser.add_argument("--temperature", type=float, default=1.0, help="sampling temperature") +parser.add_argument("--top_k", type=int, default=50, help="top-k sampling (0 = disabled)") +# Optimization +parser.add_argument("--embedding_lr", type=float, default=0.2, help="learning rate for embedding parameters (Adam)") +parser.add_argument("--unembedding_lr", type=float, default=0.004, help="learning rate for unembedding parameters (Adam)") +parser.add_argument("--matrix_lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)") +parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay for embedding/unembedding parameters (Adam)") +parser.add_argument("--init_lr_frac", type=float, default=0.05, help="initial LR as fraction of base LR") +# Evaluation / checkpointing +parser.add_argument("--eval_every", type=int, default=60, help="evaluate pass@k every N steps") +parser.add_argument("--eval_examples", type=int, default=400, help="number of examples for pass@k evaluation") +parser.add_argument("--save_every", type=int, default=60, help="save checkpoint every N steps") +args = parser.parse_args() +user_config = vars(args).copy() # ----------------------------------------------------------------------------- # Init compute/precision -ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init() +device_type = autodetect_device_type() if args.device_type == "" else args.device_type +ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. -dtype = torch.float32 if dtype == 'float32' else torch.bfloat16 -autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=dtype) +ptdtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16 +autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext() # wandb logging init -use_dummy_wandb = run == "dummy" or not master_process -wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-rl", name=run, config=user_config) +use_dummy_wandb = args.run == "dummy" or not master_process +wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-rl", name=args.run, config=user_config) # Init model and tokenizer -model, tokenizer, meta = load_model(source, device, phase="eval") +model, tokenizer, meta = load_model(args.source, device, phase="eval", model_tag=args.model_tag, step=args.model_step) engine = Engine(model, tokenizer) # for sampling rollouts # ----------------------------------------------------------------------------- @@ -72,7 +86,7 @@ engine = Engine(model, tokenizer) # for sampling rollouts train_task = GSM8K(subset="main", split="train") val_task = GSM8K(subset="main", split="test") -num_steps = (len(train_task) // examples_per_step) * num_epochs +num_steps = (len(train_task) // args.examples_per_step) * args.num_epochs print0(f"Calculated number of steps: {num_steps}") @torch.no_grad() @@ -93,16 +107,16 @@ def get_batch(): model.eval() # ensure the model is in eval mode generated_token_sequences = [] masks = [] - num_sampling_steps = num_samples // device_batch_size # go sequentially to prevent OOMs + num_sampling_steps = args.num_samples // args.device_batch_size # go sequentially to prevent OOMs for sampling_step in range(num_sampling_steps): seed = hash((step, example_idx, sampling_step)) & 0x7FFFFFFF # positive half of int32 with autocast_ctx: generated_token_sequences_batch, masks_batch = engine.generate_batch( tokens, - num_samples=device_batch_size, - max_tokens=max_new_tokens, - temperature=temperature, - top_k=top_k, + num_samples=args.device_batch_size, + max_tokens=args.max_new_tokens, + temperature=args.temperature, + top_k=args.top_k, seed=seed, # must make sure to change the seed for each sampling step ) generated_token_sequences.extend(generated_token_sequences_batch) @@ -189,16 +203,16 @@ def run_gsm8k_eval(task, tokenizer, engine, # Init the optimizer optimizers = model.setup_optimizers( - unembedding_lr=unembedding_lr, - embedding_lr=embedding_lr, - matrix_lr=matrix_lr, - weight_decay=weight_decay, + unembedding_lr=args.unembedding_lr, + embedding_lr=args.embedding_lr, + matrix_lr=args.matrix_lr, + weight_decay=args.weight_decay, ) # Set the initial learning rate as a fraction of the base learning rate for opt in optimizers: for group in opt.param_groups: - group["lr"] = group["lr"] * init_lr_frac + group["lr"] = group["lr"] * args.init_lr_frac group["initial_lr"] = group["lr"] # save the initial learning so we can decay easily later # Learning rate scheduler: simple rampdown to zero over num_steps @@ -207,9 +221,9 @@ def get_lr_multiplier(it): return lrm # Calculate the number of examples each rank handles to achieve the desired examples_per_step -print0(f"Total sequences per step: {examples_per_step * num_samples}") # total batch size in sequences/step -assert examples_per_step % ddp_world_size == 0, "Desired examples per step must be divisible by the number of ranks" -examples_per_rank = examples_per_step // ddp_world_size # per GPU +print0(f"Total sequences per step: {args.examples_per_step * args.num_samples}") # total batch size in sequences/step +assert args.examples_per_step % ddp_world_size == 0, "Desired examples per step must be divisible by the number of ranks" +examples_per_rank = args.examples_per_step // ddp_world_size # per GPU print0(f"Calculated examples per rank: {examples_per_rank}") # Kick off the training loop @@ -217,22 +231,22 @@ batch_iterator = get_batch() for step in range(num_steps): # Evaluate the model once in a while and log to wandb - if step % eval_every == 0: + if step % args.eval_every == 0: model.eval() - passk = torch.zeros(device_batch_size, device=device) # pass@k for k=1..device_batch_size + passk = torch.zeros(args.device_batch_size, device=device) # pass@k for k=1..device_batch_size with autocast_ctx: - records_iter = run_gsm8k_eval(val_task, tokenizer, engine, num_samples=device_batch_size, max_examples=eval_examples, temperature=1.0) + records_iter = run_gsm8k_eval(val_task, tokenizer, engine, num_samples=args.device_batch_size, max_examples=args.eval_examples, temperature=1.0) records = list(records_iter) # collect all records - for k in range(1, device_batch_size + 1): + for k in range(1, args.device_batch_size + 1): passk[k - 1] = sum(any(o["is_correct"] for o in r["outcomes"][:k]) for r in records) num_records = torch.tensor(len(records), dtype=torch.long, device=device) if ddp: dist.all_reduce(num_records, op=dist.ReduceOp.SUM) dist.all_reduce(passk, op=dist.ReduceOp.SUM) passk = passk / num_records.item() # normalize by the total number of records - print_passk = [f"Pass@{k}: {passk[k - 1].item():.4f}" for k in range(1, device_batch_size + 1)] + print_passk = [f"Pass@{k}: {passk[k - 1].item():.4f}" for k in range(1, args.device_batch_size + 1)] print0(f"Step {step} | {', '.join(print_passk)}") - log_passk = {f"pass@{k}": passk[k - 1].item() for k in range(1, device_batch_size + 1)} + log_passk = {f"pass@{k}": passk[k - 1].item() for k in range(1, args.device_batch_size + 1)} wandb_run.log({ "step": step, **log_passk, @@ -247,11 +261,11 @@ for step in range(num_steps): # Evaluate the loss and gradients model.train() # ensure the model is in train mode # We need one more loop because we can never exceed the device_batch_size - assert inputs_all.size(0) % device_batch_size == 0 - num_passes = inputs_all.size(0) // device_batch_size + assert inputs_all.size(0) % args.device_batch_size == 0 + num_passes = inputs_all.size(0) // args.device_batch_size for pass_idx in range(num_passes): # Pluck out the batch for this pass - b0, b1 = pass_idx * device_batch_size, (pass_idx + 1) * device_batch_size + b0, b1 = pass_idx * args.device_batch_size, (pass_idx + 1) * args.device_batch_size inputs = inputs_all[b0:b1] targets = targets_all[b0:b1] rewards = rewards_all[b0:b1] @@ -304,11 +318,11 @@ for step in range(num_steps): }) # Master process saves the model once in a while. Skip first step. Save last step. - if master_process and ((step > 0 and step % save_every == 0) or step == num_steps - 1): + if master_process and ((step > 0 and step % args.save_every == 0) or step == num_steps - 1): base_dir = get_base_dir() depth = model.config.n_layer - model_tag = f"d{depth}" # base the model tag on the depth of the base model - checkpoint_dir = os.path.join(base_dir, "chatrl_checkpoints", model_tag) + output_dirname = args.model_tag if args.model_tag else f"d{depth}" # base the model tag on the depth of the base model + checkpoint_dir = os.path.join(base_dir, "chatrl_checkpoints", output_dirname) model_config_kwargs = model.config.__dict__ # slightly naughty, abusing the simplicity of GPTConfig, TODO nicer save_checkpoint( checkpoint_dir, diff --git a/scripts/chat_sft.py b/scripts/chat_sft.py index bbeb1f9d..853a2bf6 100644 --- a/scripts/chat_sft.py +++ b/scripts/chat_sft.py @@ -9,8 +9,9 @@ Or torchrun for training: torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft """ +import argparse import os -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" +os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" import wandb import torch @@ -31,49 +32,51 @@ from tasks.customjson import CustomJSON from tasks.spellingbee import SimpleSpelling, SpellingBee # ----------------------------------------------------------------------------- -# SFT Hyperparameters -run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb) -# input model options -source = "mid" # base|mid , which checkpoint to load the model from (base model or midtrained model) -model_tag = None # model tag to load the model from (base model or midtrained model) -step = None # step to load the model from (base model or midtrained model) -# compute/precision -device_type = "" # cuda|cpu|mps (empty => autodetect) -dtype = "bfloat16" -device_batch_size = 4 # max to avoid OOM -# optimization -num_epochs = 1 -num_iterations = -1 # override number of iterations (-1 = disable, use num_epochs to derive it) -target_examples_per_step = 32 -unembedding_lr = 0.004 -embedding_lr = 0.2 -matrix_lr = 0.02 -weight_decay = 0.0 -init_lr_frac = 0.02 -# evaluation and logging there of -eval_every = 100 -eval_steps = 100 -eval_metrics_every = 200 -eval_metrics_max_problems = 1024 -# now allow CLI to override the settings via the configurator lol -config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))] -exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file -user_config = {k: globals()[k] for k in config_keys} # possibly useful for logging +# CLI arguments +parser = argparse.ArgumentParser(description="Supervised finetuning for chat") +# Logging +parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") +# Runtime +parser.add_argument("--device_type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") +parser.add_argument("--dtype", type=str, default="bfloat16", help="float32|bfloat16") +# Model loading +parser.add_argument("--source", type=str, default="mid", help="base|mid - which checkpoint to load from") +parser.add_argument("--model_tag", type=str, default=None, help="model tag to load from") +parser.add_argument("--model_step", type=int, default=None, help="model step to load from") +# Training horizon +parser.add_argument("--num_epochs", type=int, default=1, help="number of epochs") +parser.add_argument("--num_iterations", type=int, default=-1, help="override number of iterations (-1 = use num_epochs)") +# Batch sizes +parser.add_argument("--device_batch_size", type=int, default=4, help="per-device batch size") +parser.add_argument("--target_examples_per_step", type=int, default=32, help="target examples per optimization step") +# Optimization +parser.add_argument("--embedding_lr", type=float, default=0.2, help="learning rate for embedding parameters (Adam)") +parser.add_argument("--unembedding_lr", type=float, default=0.004, help="learning rate for unembedding parameters (Adam)") +parser.add_argument("--matrix_lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)") +parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay for embedding/unembedding parameters (Adam)") +parser.add_argument("--init_lr_frac", type=float, default=0.02, help="initial LR as fraction of base LR") +# Evaluation +parser.add_argument("--eval_every", type=int, default=100, help="evaluate val loss every N steps") +parser.add_argument("--eval_steps", type=int, default=100, help="number of batches for val loss evaluation") +parser.add_argument("--eval_metrics_every", type=int, default=200, help="evaluate accuracy metrics every N steps") +parser.add_argument("--eval_metrics_max_problems", type=int, default=1024, help="max problems per metric evaluation") +args = parser.parse_args() +user_config = vars(args).copy() # ----------------------------------------------------------------------------- # Compute init -device_type = autodetect_device_type() if device_type == "" else device_type +device_type = autodetect_device_type() if args.device_type == "" else args.device_type ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) master_process = ddp_rank == 0 -ptdtype = torch.float32 if dtype == 'float32' else torch.bfloat16 +ptdtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16 autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext() # wandb logging init -use_dummy_wandb = run == "dummy" or not master_process -wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-sft", name=run, config=user_config, save_code=True) +use_dummy_wandb = args.run == "dummy" or not master_process +wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-sft", name=args.run, config=user_config, save_code=True) # Load the model and tokenizer -model, tokenizer, meta = load_model(source, device, phase="train", model_tag=model_tag, step=step) +model, tokenizer, meta = load_model(args.source, device, phase="train", model_tag=args.model_tag, step=args.model_step) orig_model = model # original, uncompiled model # model = torch.compile(model, dynamic=True) # doesn't work super well because of variable lengths of inputs engine = Engine(model, tokenizer) # will be used for inline model evaluation only @@ -127,34 +130,36 @@ def sft_data_generator(dataset, batch_size): yield collate_and_yield(batch) batch = [] -examples_per_step = device_batch_size * ddp_world_size -print0(f"Target examples per step: {target_examples_per_step}") -print0(f"Device batch size: {device_batch_size}") +examples_per_step = args.device_batch_size * ddp_world_size +print0(f"Target examples per step: {args.target_examples_per_step}") +print0(f"Device batch size: {args.device_batch_size}") print0(f"Examples per step is device_batch_size * ddp_world_size: {examples_per_step}") -assert target_examples_per_step % examples_per_step == 0, "Target examples per step must be divisible by examples per step" -grad_accum_steps = target_examples_per_step // examples_per_step +assert args.target_examples_per_step % examples_per_step == 0, "Target examples per step must be divisible by examples per step" +grad_accum_steps = args.target_examples_per_step // examples_per_step print0(f"=> Setting grad accum steps: {grad_accum_steps}") -if num_iterations == -1: +if args.num_iterations == -1: # derive num_iterations from num_epochs and the size of the dataset - assert num_epochs > 0, "num_epochs must be positive if num_iterations is -1" - num_iterations = (len(train_ds) // target_examples_per_step) * num_epochs -train_loader = sft_data_generator(train_ds, batch_size=device_batch_size) -build_val_loader = lambda: sft_data_generator(val_ds, batch_size=device_batch_size) + assert args.num_epochs > 0, "num_epochs must be positive if num_iterations is -1" + num_iterations = (len(train_ds) // args.target_examples_per_step) * args.num_epochs +else: + num_iterations = args.num_iterations +train_loader = sft_data_generator(train_ds, batch_size=args.device_batch_size) +build_val_loader = lambda: sft_data_generator(val_ds, batch_size=args.device_batch_size) # ----------------------------------------------------------------------------- # Initialize the Optimizer optimizers = model.setup_optimizers( - unembedding_lr=unembedding_lr, - embedding_lr=embedding_lr, - matrix_lr=matrix_lr, - weight_decay=weight_decay, + unembedding_lr=args.unembedding_lr, + embedding_lr=args.embedding_lr, + matrix_lr=args.matrix_lr, + weight_decay=args.weight_decay, ) # Set the initial learning rate as a fraction of the base learning rate for opt in optimizers: for group in opt.param_groups: - group["lr"] = group["lr"] * init_lr_frac + group["lr"] = group["lr"] * args.init_lr_frac group["initial_lr"] = group["lr"] # save the initial learning so we can decay easily later # ----------------------------------------------------------------------------- @@ -167,17 +172,16 @@ def get_lr_multiplier(it): # Go! step = 0 -train_iter = iter(train_loader) for step in range(num_iterations): last_step = step == num_iterations - 1 # evaluate the validation loss - if last_step or step % eval_every == 0: + if last_step or step % args.eval_every == 0: model.eval() - val_iter = iter(build_val_loader()) + val_loader = build_val_loader() losses = [] - for _ in range(eval_steps): - val_inputs, val_targets = next(val_iter) + for _ in range(args.eval_steps): + val_inputs, val_targets = next(val_loader) with torch.no_grad(), autocast_ctx: loss = model(val_inputs, val_targets) losses.append(loss) @@ -193,13 +197,13 @@ for step in range(num_iterations): model.train() # evaluate accuracy of the multiple choice tasks (which are quick to run) - if last_step or (step > 0 and step % eval_metrics_every == 0): + if last_step or (step > 0 and step % args.eval_metrics_every == 0): model.eval() metrics = {} with torch.no_grad(), autocast_ctx: # note that because these are inside no_grad, we can usually afford to at least ~2X the batch size - metrics["mmlu_acc"] = run_chat_eval("MMLU", model, tokenizer, engine, batch_size=device_batch_size*2, max_problems=eval_metrics_max_problems) - metrics["arc_easy_acc"] = run_chat_eval("ARC-Easy", model, tokenizer, engine, batch_size=device_batch_size*2, max_problems=eval_metrics_max_problems) + metrics["mmlu_acc"] = run_chat_eval("MMLU", model, tokenizer, engine, batch_size=args.device_batch_size*2, max_problems=args.eval_metrics_max_problems) + metrics["arc_easy_acc"] = run_chat_eval("ARC-Easy", model, tokenizer, engine, batch_size=args.device_batch_size*2, max_problems=args.eval_metrics_max_problems) metrics_str = ', '.join(f'{k}: {v:.6f}' for k, v in metrics.items()) print0(f"Step {step:05d} | {metrics_str}") wandb_run.log({ @@ -214,7 +218,7 @@ for step in range(num_iterations): # evaluate the gradient num_tokens = torch.tensor(0, device=device) # the number of "active" tokens of supervision seen for micro_step in range(grad_accum_steps): - train_inputs, train_targets = next(train_iter) + train_inputs, train_targets = next(train_loader) with autocast_ctx: loss = model(train_inputs, train_targets) train_loss = loss.detach() # for logging @@ -251,8 +255,8 @@ for step in range(num_iterations): if master_process: base_dir = get_base_dir() depth = model.config.n_layer - model_tag = f"d{depth}" # base the model tag on the depth of the base model - checkpoint_dir = os.path.join(base_dir, "chatsft_checkpoints", model_tag) + output_dirname = args.model_tag if args.model_tag else f"d{depth}" # e.g. d12 + checkpoint_dir = os.path.join(base_dir, "chatsft_checkpoints", output_dirname) model_config_kwargs = model.config.__dict__ # slightly naughty, abusing the simplicity of GPTConfig, TODO nicer save_checkpoint( checkpoint_dir, diff --git a/scripts/mid_train.py b/scripts/mid_train.py index aa577247..ba90ccb0 100644 --- a/scripts/mid_train.py +++ b/scripts/mid_train.py @@ -9,9 +9,10 @@ Or torchrun for training: torchrun --standalone --nproc_per_node=8 -m scripts.mid_train -- --device_batch_size=16 """ +import argparse from collections import deque import os -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" +os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" import time import wandb import torch @@ -31,65 +32,75 @@ from tasks.customjson import CustomJSON from tasks.spellingbee import SimpleSpelling, SpellingBee # ----------------------------------------------------------------------------- -run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb) -device_type = "" # cuda|cpu|mps (empty => autodetect) -model_tag = None # model tag to load the model from (base model or midtrained model) -step = None # step to load the model from (base model or midtrained model) -dtype = "bfloat16" -num_iterations = -1 # explicit number of steps of the optimization (-1 = disable) -max_seq_len = 2048 -device_batch_size = 32 -unembedding_lr = 0.004 -embedding_lr = 0.2 -matrix_lr = 0.02 -init_lr_frac = 1.0 # initial learning rate is this fraction of the base learning rate -weight_decay = 0.0 -eval_every = 150 # -1 = disable -eval_tokens = 20*524288 -total_batch_size = 524288 -dry_run = 0 # dry_run=1 is for experiments: we will log to wandb but we won't write checkpoints or report -config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))] -exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file -user_config = {k: globals()[k] for k in config_keys} # possibly useful for logging +# CLI arguments +parser = argparse.ArgumentParser(description="Midtrain the model") +# Logging +parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") +# Runtime +parser.add_argument("--device_type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") +parser.add_argument("--dtype", type=str, default="bfloat16", help="float32|bfloat16") +# Model loading +parser.add_argument("--model_tag", type=str, default=None, help="model tag to load from") +parser.add_argument("--model_step", type=int, default=None, help="model step to load from") +# Training horizon +parser.add_argument("--num_iterations", type=int, default=-1, help="number of optimization steps (-1 = full epoch)") +# Batch sizes +parser.add_argument("--max_seq_len", type=int, default=2048, help="max context length") +parser.add_argument("--device_batch_size", type=int, default=32, help="per-device batch size") +parser.add_argument("--total_batch_size", type=int, default=524288, help="total batch size in tokens") +# Optimization +parser.add_argument("--embedding_lr", type=float, default=0.2, help="learning rate for embedding parameters (Adam)") +parser.add_argument("--unembedding_lr", type=float, default=0.004, help="learning rate for unembedding parameters (Adam)") +parser.add_argument("--matrix_lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)") +parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay for embedding/unembedding parameters (Adam)") +parser.add_argument("--init_lr_frac", type=float, default=1.0, help="initial LR as fraction of base LR") +# Evaluation +parser.add_argument("--eval_every", type=int, default=150, help="evaluate val bpb every N steps (-1 = disable)") +parser.add_argument("--eval_tokens", type=int, default=20*524288, help="number of tokens to evaluate val loss on") +# Output +parser.add_argument("--dry_run", action="store_true", help="log to wandb but skip checkpoints/report") +args = parser.parse_args() +user_config = vars(args).copy() # ----------------------------------------------------------------------------- # Compute init -device_type = autodetect_device_type() if device_type == "" else device_type +device_type = autodetect_device_type() if args.device_type == "" else args.device_type ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) master_process = ddp_rank == 0 -autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext() +ptdtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16 +autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext() synchronize = torch.cuda.synchronize if device_type == "cuda" else lambda: None get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0 # wandb logging init -use_dummy_wandb = run == "dummy" or not master_process -wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-mid", name=run, config=user_config) +use_dummy_wandb = args.run == "dummy" or not master_process +wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-mid", name=args.run, config=user_config) # Load the model and tokenizer -model, tokenizer, meta = load_model("base", device, phase="train", model_tag=model_tag, step=step) +model, tokenizer, meta = load_model("base", device, phase="train", model_tag=args.model_tag, step=args.model_step) pretrain_batch_size = meta.get("device_batch_size", None) -if pretrain_batch_size is not None and device_batch_size > pretrain_batch_size: +if pretrain_batch_size is not None and args.device_batch_size > pretrain_batch_size: print0(f"FOOTGUN WARNING: base model training used device_batch_size {pretrain_batch_size}, did you pass in a good --device_batch_size to this script?") orig_model = model model = torch.compile(model, dynamic=False) depth = model.config.n_layer num_flops_per_token = model.estimate_flops() -tokens_per_fwdbwd = device_batch_size * max_seq_len # tokens per iteration for a single rank +tokens_per_fwdbwd = args.device_batch_size * args.max_seq_len # tokens per iteration for a single rank world_tokens_per_fwdbwd = tokens_per_fwdbwd * ddp_world_size # total tokens per iteration for all ranks -assert total_batch_size % world_tokens_per_fwdbwd == 0 -grad_accum_steps = total_batch_size // world_tokens_per_fwdbwd -print0(f"Tokens / micro-batch / rank: {device_batch_size} x {max_seq_len} = {tokens_per_fwdbwd:,}") +assert args.total_batch_size % world_tokens_per_fwdbwd == 0 +grad_accum_steps = args.total_batch_size // world_tokens_per_fwdbwd +print0(f"Tokens / micro-batch / rank: {args.device_batch_size} x {args.max_seq_len} = {tokens_per_fwdbwd:,}") print0(f"Tokens / micro-batch: {world_tokens_per_fwdbwd:,}") -print0(f"Total batch size {total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}") +print0(f"Total batch size {args.total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}") token_bytes = get_token_bytes(device=device) # Initialize the Optimizer (Muon for Linear layers, AdamW for embedding and lm_head) -optimizers = model.setup_optimizers(unembedding_lr=unembedding_lr, embedding_lr=embedding_lr, matrix_lr=matrix_lr, weight_decay=weight_decay) +optimizers = model.setup_optimizers(unembedding_lr=args.unembedding_lr, embedding_lr=args.embedding_lr, matrix_lr=args.matrix_lr, weight_decay=args.weight_decay) adamw_optimizer, muon_optimizer = optimizers # Override the initial learning rate as a fraction of the base learning rate for opt in optimizers: for group in opt.param_groups: - group["lr"] = group["lr"] * init_lr_frac + group["lr"] = group["lr"] * args.init_lr_frac group["initial_lr"] = group["lr"] # save the initial learning so we can decay easily later # Midtraining data mixture and DataLoader @@ -112,7 +123,7 @@ val_dataset = TaskMixture([ # DataLoader is defined here, it emits inputs, targets : 2D tensors of shape (device_batch_size, max_seq_len) # A big problem is that we don't know the final num_iterations in advance. So we create # these two global variables and update them from within the data generator. -last_step = False # we will toggle this to True when we reach the end of the dataset +last_step = False # we will toggle this to True when we reach the end of the training dataset approx_progress = 0.0 # will go from 0 to 1 over the course of the epoch def mid_data_generator(split): global last_step, approx_progress @@ -120,7 +131,7 @@ def mid_data_generator(split): dataset = train_dataset if split == "train" else val_dataset dataset_size = len(dataset) assert dataset_size > 0 - needed_tokens = device_batch_size * max_seq_len + 1 # to form one training batch of inputs,targets + needed_tokens = args.device_batch_size * args.max_seq_len + 1 # to form one training batch of inputs,targets token_buffer = deque() # CUDA supports memory pinning for faster transfers between CPU and GPU: scratch = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=(device_type == "cuda")) @@ -139,18 +150,18 @@ def mid_data_generator(split): last_step = True # toggle last_step to True, which will terminate the training loop # Stopping condition to respect num_iterations, if given it += 1 - if num_iterations > 0 and it >= num_iterations: + if 0 < args.num_iterations <= it and split == "train": last_step = True # toggle last_step to True, which will terminate the training loop # Build up inputs/targets and yield for i in range(needed_tokens): scratch[i] = token_buffer.popleft() inputs_cpu = scratch[:-1].to(dtype=torch.int32) targets_cpu = scratch[1:] - inputs = inputs_cpu.view(device_batch_size, max_seq_len).to(device=device, dtype=torch.int32, non_blocking=True) - targets = targets_cpu.view(device_batch_size, max_seq_len).to(device=device, dtype=torch.int64, non_blocking=True) + inputs = inputs_cpu.view(args.device_batch_size, args.max_seq_len).to(device=device, dtype=torch.int32, non_blocking=True) + targets = targets_cpu.view(args.device_batch_size, args.max_seq_len).to(device=device, dtype=torch.int64, non_blocking=True) if split == "train": - if num_iterations > 0: - approx_progress = it / num_iterations # calculate progress from the max number of iterations + if args.num_iterations > 0: + approx_progress = it / args.num_iterations # calculate progress from the max number of iterations else: approx_progress = cursor / dataset_size # approximate progress as a fraction of the dataset yield inputs, targets @@ -179,7 +190,7 @@ ema_beta = 0.9 # EMA decay factor total_training_time = 0 # total wall-clock time of training step = 0 while True: - flops_so_far = num_flops_per_token * total_batch_size * step + flops_so_far = num_flops_per_token * args.total_batch_size * step # Synchronize last_step across all ranks to avoid hangs in the distributed setting if ddp: @@ -188,10 +199,10 @@ while True: last_step = bool(last_step_tensor.item()) # once in a while: evaluate the val bpb (all ranks participate) - if last_step or (eval_every > 0 and step % eval_every == 0): + if last_step or (args.eval_every > 0 and step % args.eval_every == 0): model.eval() val_loader = build_val_loader() - eval_steps = eval_tokens // (device_batch_size * max_seq_len * ddp_world_size) + eval_steps = args.eval_tokens // (args.device_batch_size * args.max_seq_len * ddp_world_size) with autocast_ctx: val_bpb = evaluate_bpb(model, val_loader, eval_steps, token_bytes) print0(f"Step {step:05d} | Validation bpb: {val_bpb:.4f}") @@ -206,8 +217,8 @@ while True: model.train() # save checkpoint at the end of the run (only on master process) - if master_process and last_step and not dry_run: - output_dirname = f"d{depth}" # e.g. d12 + if master_process and last_step and not args.dry_run: + output_dirname = args.model_tag if args.model_tag else f"d{depth}" # e.g. d12 checkpoint_dir = os.path.join(base_dir, "mid_checkpoints", output_dirname) save_checkpoint( checkpoint_dir, @@ -218,7 +229,7 @@ while True: "step": step, "val_bpb": val_bpb, # loss at last step "model_config": { - "sequence_len": max_seq_len, + "sequence_len": args.max_seq_len, "vocab_size": tokenizer.get_vocab_size(), "n_layer": depth, "n_head": model.config.n_head, @@ -268,8 +279,8 @@ while True: smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss.item() # EMA the training loss debiased_smooth_loss = smooth_train_loss / (1 - ema_beta**(step + 1)) # debias the EMA pct_done = 100 * progress - tok_per_sec = int(total_batch_size / dt) - flops_per_sec = num_flops_per_token * total_batch_size / dt + tok_per_sec = int(args.total_batch_size / dt) + flops_per_sec = num_flops_per_token * args.total_batch_size / dt promised_flops_per_sec_h100 = 989e12 * ddp_world_size # bfloat16 H100 SXM and without 2:4 sparsity mfu = 100 * flops_per_sec / promised_flops_per_sec_h100 # in % if step > 10: @@ -293,7 +304,7 @@ print0(f"Total training time: {total_training_time/60:.2f}m") print0(f"Minimum validation bpb: {min_val_bpb:.4f}") # Log to report -if not dry_run: +if not args.dry_run: from nanochat.report import get_report get_report().log(section="Midtraining", data=[ user_config, # CLI args diff --git a/scripts/tok_train.py b/scripts/tok_train.py index c2faf17e..e1b79ee2 100644 --- a/scripts/tok_train.py +++ b/scripts/tok_train.py @@ -1,5 +1,5 @@ """ -Train a tokenizer using the HuggingFace Tokenizers library. +Train a tokenizer using our own BPE Tokenizer library. In the style of GPT-4 tokenizer. """ import os diff --git a/speedrun.sh b/speedrun.sh index 7955ec51..8803dcb4 100644 --- a/speedrun.sh +++ b/speedrun.sh @@ -48,13 +48,6 @@ python -m nanochat.report reset # ----------------------------------------------------------------------------- # Tokenizer -# Install Rust / Cargo -curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y -source "$HOME/.cargo/env" - -# Build the rustbpe Tokenizer -uv run maturin develop --release --manifest-path rustbpe/Cargo.toml - # Download the first ~2B characters of pretraining dataset # look at dev/repackage_data_reference.py for details on how this data was prepared # each data shard is ~250M chars @@ -96,7 +89,7 @@ torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_eval # Midtraining (teach the model conversation special tokens, tool use, multiple choice) # download 2.3MB of synthetic identity conversations to impart a personality to nanochat -# see dev/gen_sft_data.py for details on how this data was prepared and to get a sense of how you can easily tune it +# see dev/gen_synthetic_data.py for details on how this data was prepared and to get a sense of how you can easily tune it curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl # run midtraining and eval the model diff --git a/tasks/common.py b/tasks/common.py index dcd2e910..2d6ddd86 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -53,7 +53,7 @@ class Task: class TaskMixture(Task): """ - For SFT Training it becomes useful to train on a tax mixture of datasets. + For SFT Training it becomes useful to train on a mixture of datasets. Fun trick: if you wish to oversample any task, just pass it in multiple times in the list. """ diff --git a/tasks/gsm8k.py b/tasks/gsm8k.py index c05e21cd..f3539ff0 100644 --- a/tasks/gsm8k.py +++ b/tasks/gsm8k.py @@ -74,7 +74,7 @@ class GSM8K(Task): else: # Regular text in between tool calls assistant_message_parts.append({"type": "text", "text": part}) - # No put it all together + # Now put it all together messages = [ {"role": "user", "content": question}, # note: simple string {"role": "assistant", "content": assistant_message_parts}, # note: list of parts (as dicts) diff --git a/tasks/spellingbee.py b/tasks/spellingbee.py index 3b453055..24954c00 100644 --- a/tasks/spellingbee.py +++ b/tasks/spellingbee.py @@ -35,6 +35,8 @@ from nanochat.common import download_file_with_lock LETTERS = "abcdefghijklmnopqrstuvwxyz" # A list of 370K English words of large variety WORD_LIST_URL = "https://raw.githubusercontent.com/dwyl/english-words/refs/heads/master/words_alpha.txt" +# A number bigger than 370K to separate train and test random seeds +TEST_RANDOM_SEED_OFFSET = 10_000_000 # Identical to gsm8k's answer extraction ANSWER_RE = re.compile(r"#### (\-?[0-9\.\,]+)") @@ -131,7 +133,7 @@ class SpellingBee(Task): return self.size def get_example(self, index): - seed = index if self.split == "train" else -(index + 1) # avoid collision at 0 + seed = index if self.split == 'train' else TEST_RANDOM_SEED_OFFSET + index rng = random.Random(seed) # pick a random word @@ -252,7 +254,7 @@ class SimpleSpelling(Task): return self.size def get_example(self, index): - seed = index if self.split == "train" else -(index + 1) # avoid collision at 0 + seed = index if self.split == 'train' else TEST_RANDOM_SEED_OFFSET + index rng = random.Random(seed) # pick a random word word = rng.choice(self.words) diff --git a/tests/test_engine.py b/tests/test_engine.py index 7403b36a..683f89bf 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -5,7 +5,85 @@ python -m pytest tests/test_engine.py -v """ import torch -from nanochat.engine import KVCache +from nanochat.engine import KVCache, Engine +from dataclasses import dataclass + + +# ----------------------------------------------------------------------------- +# Mock classes for testing Engine without loading a real model + +@dataclass +class MockConfig: + """Minimal config for Engine tests.""" + n_kv_head: int = 4 + n_head: int = 4 + n_embd: int = 64 + n_layer: int = 2 + sequence_len: int = 128 + + +class MockModel: + """ + Mock model that returns uniform logits over the vocab. + This ensures that with temperature > 0, different samples should + (with very high probability) produce different tokens. + """ + def __init__(self, vocab_size=262): # 256 bytes + 6 special tokens + self.vocab_size = vocab_size + self.config = MockConfig() + self._device = "cpu" + + def get_device(self): + return self._device + + def forward(self, ids, kv_cache=None): + """Return uniform logits so sampling is spread across vocab.""" + B, T = ids.shape + # Simulate what a real transformer does: insert k,v into the cache for each layer + if kv_cache is not None: + head_dim = self.config.n_embd // self.config.n_head + for layer_idx in range(self.config.n_layer): + k = torch.zeros(B, self.config.n_kv_head, T, head_dim) + v = torch.zeros(B, self.config.n_kv_head, T, head_dim) + kv_cache.insert_kv(layer_idx, k, v) + # Uniform logits -> equal probability for all tokens + logits = torch.zeros(B, T, self.vocab_size) + return logits + + +class ByteTokenizer: + """ + Simple byte-level tokenizer for testing. + Tokens 0-255 are raw bytes, 256+ are special tokens. + """ + def __init__(self): + # Special tokens start at 256 + self._special_tokens = { + "<|python_start|>": 256, + "<|python_end|>": 257, + "<|output_start|>": 258, + "<|output_end|>": 259, + "<|assistant_end|>": 260, + "<|bos|>": 261, + } + self._bos = 261 + + def encode_special(self, s): + return self._special_tokens[s] + + def get_bos_token_id(self): + return self._bos + + def encode(self, s, prepend=None): + tokens = list(s.encode("utf-8")) # bytes 0-255 + if prepend is not None: + tokens = [prepend] + tokens + return tokens + + def decode(self, tokens): + # Filter out special tokens before decoding + byte_tokens = [t for t in tokens if t < 256] + return bytes(byte_tokens).decode("utf-8", errors="replace") def test_kv_cache_resize(): """ @@ -64,3 +142,46 @@ def test_kv_cache_resize(): original_v = original_cache[layer_idx, 1, :, :, token_idx, :] assert (actual_k == original_k).all(), f"Layer {layer_idx}, token {token_idx}: key doesn't match original" assert (actual_v == original_v).all(), f"Layer {layer_idx}, token {token_idx}: value doesn't match original" + + +def test_multi_sample_first_token_diversity(): + """ + Test that when generating multiple samples, each sample gets an independently + sampled first token (not a broadcast of the same token to all rows). + + Previously, the first token after prefill was sampled once and broadcast to all + rows, causing all samples to start identically. The fix expands the prefill logits + to num_samples and samples independently for each row. + + With uniform logits over 262 tokens and 16 samples, the probability that all + samples independently pick the same token is (1/262)^15 ≈ 10^-36. So if they're + all identical, it indicates tokens are being broadcast instead of independently sampled. + """ + model = MockModel(vocab_size=262) + tokenizer = ByteTokenizer() + engine = Engine(model, tokenizer) + + # Generate 16 samples with temperature=1.0 (stochastic sampling) + prompt_tokens = [261, 72, 101, 108, 108, 111] # + "Hello" + num_samples = 16 + + # Collect the first generated token from each sample + first_tokens = [] + gen = engine.generate( + prompt_tokens, + num_samples=num_samples, + max_tokens=1, # We only need the first token + temperature=1.0, + seed=42, + ) + for token_column, token_masks in gen: + first_tokens = token_column # This is the first (and only) yield + + # With uniform distribution and 16 samples, they should NOT all be identical + # If they are all identical, the bug exists (broadcasting instead of sampling) + unique_tokens = set(first_tokens) + assert len(unique_tokens) > 1, ( + f"All {num_samples} samples got the same first token ({first_tokens[0]}). " + f"With uniform logits, this is statistically impossible (~10^-36 probability) " + f"unless tokens are being broadcast instead of independently sampled." + ) diff --git a/tests/test_rustbpe.py b/tests/test_rustbpe.py deleted file mode 100644 index aca67fcd..00000000 --- a/tests/test_rustbpe.py +++ /dev/null @@ -1,635 +0,0 @@ -""" -Comparing the training of: - -1. (very slow) Python reference implementation -2. Optimized Python implementation -3. HuggingFace tokenizers training implementation -4. Our own custom RustBPE training implementation - -All of these should calculate the same merges and produce -the same vocabulary and tokenizations. - -Finally, for inference we will use tiktoken for efficiency. -So we want to make sure we can export our rustbpe tokenizer -into tiktoken and use it for inference with identical results. - -Run with: -python -m pytest tests/test_rustbpe.py -v -s --v is verbose, -s is show prints -""" - -import regex as re -from collections import Counter, defaultdict -import time -import rustbpe -import tiktoken -import pytest - -GPT4_SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""" - -# ----------------------------------------------------------------------------- -# Reference tokenizer, pretty much copy pasted and pruned a bit from minbpe - -def get_stats(ids, counts=None): - """ - Given a list of integers, return a dictionary of counts of consecutive pairs - Example: [1, 2, 3, 1, 2] -> {(1, 2): 2, (2, 3): 1, (3, 1): 1} - Optionally allows to update an existing dictionary of counts - """ - counts = {} if counts is None else counts - for pair in zip(ids, ids[1:]): # iterate consecutive elements - counts[pair] = counts.get(pair, 0) + 1 - return counts - -def merge(ids, pair, idx): - """ - In the list of integers (ids), replace all consecutive occurrences - of pair with the new integer token idx - Example: ids=[1, 2, 3, 1, 2], pair=(1, 2), idx=4 -> [4, 3, 4] - """ - newids = [] - i = 0 - while i < len(ids): - # if not at the very last position AND the pair matches, replace it - if ids[i] == pair[0] and i < len(ids) - 1 and ids[i+1] == pair[1]: - newids.append(idx) - i += 2 - else: - newids.append(ids[i]) - i += 1 - return newids - -class RegexTokenizer: - - def __init__(self, pattern=None): - """ - - pattern: optional string to override the default (GPT-4 split pattern) - - special_tokens: str -> int dictionary of special tokens - example: {'<|endoftext|>': 100257} - """ - self.pattern = GPT4_SPLIT_PATTERN if pattern is None else pattern - self.merges = {} # (int, int) -> int - self.compiled_pattern = re.compile(self.pattern) - self.special_tokens = {} - self.inverse_special_tokens = {} - self.vocab = self._build_vocab() - - def _build_vocab(self): - # vocab is simply and deterministically derived from merges - vocab = {idx: bytes([idx]) for idx in range(256)} - for (p0, p1), idx in self.merges.items(): - vocab[idx] = vocab[p0] + vocab[p1] - for special, idx in self.special_tokens.items(): - vocab[idx] = special.encode("utf-8") - return vocab - - def train(self, text, vocab_size, verbose=False): - assert vocab_size >= 256 - num_merges = vocab_size - 256 - - # keep track of whether at any point during training the merge is ambiguous (counts of pairs are not unique) - ambiguous = False - - # split the text up into text chunks - text_chunks = re.findall(self.compiled_pattern, text) - - # input text preprocessing - ids = [list(ch.encode("utf-8")) for ch in text_chunks] - - # iteratively merge the most common pairs to create new tokens - merges = {} # (int, int) -> int - vocab = {idx: bytes([idx]) for idx in range(256)} # idx -> bytes - for i in range(num_merges): - # count the number of times every consecutive pair appears - stats = {} - for chunk_ids in ids: - # passing in stats will update it in place, adding up counts - get_stats(chunk_ids, stats) - # find the pair with the highest count - pair = max(stats, key=stats.get) - # check if the merge is ambiguous - i.e. the max value is not unique - pair_count = stats[pair] - pairs_with_max_count = [pair for pair, count in stats.items() if count == pair_count] - if len(pairs_with_max_count) > 1: - # print the top 10 pairs with their counts - # print(f"{i} Merge is ambiguous! {pair} has {pair_count} occurrences") - # for print_pair, print_count in sorted(stats.items(), key=lambda x: x[1], reverse=True)[:10]: - # print(f"{print_pair}: {print_count}") - ambiguous = True - # mint a new token: assign it the next available id - idx = 256 + i - # replace all occurrences of pair in ids with idx - ids = [merge(chunk_ids, pair, idx) for chunk_ids in ids] - # save the merge - merges[pair] = idx - vocab[idx] = vocab[pair[0]] + vocab[pair[1]] - # prints - if verbose: - print(f"merge {i+1}/{num_merges}: {pair} -> {idx} ({vocab[idx]}) had {stats[pair]} occurrences") - - # save class variables - self.merges = merges # used in encode() - self.vocab = vocab # used in decode() - return ambiguous - - def _encode_chunk(self, text_bytes): - # return the token ids - # let's begin. first, convert all bytes to integers in range 0..255 - ids = list(text_bytes) - while len(ids) >= 2: - # find the pair with the lowest merge index - stats = get_stats(ids) - pair = min(stats, key=lambda p: self.merges.get(p, float("inf"))) - # subtle: if there are no more merges available, the key will - # result in an inf for every single pair, and the min will be - # just the first pair in the list, arbitrarily - # we can detect this terminating case by a membership check - if pair not in self.merges: - break # nothing else can be merged anymore - # otherwise let's merge the best pair (lowest merge index) - idx = self.merges[pair] - ids = merge(ids, pair, idx) - return ids - - def encode_ordinary(self, text): - """Encoding that ignores any special tokens.""" - # split text into chunks of text by categories defined in regex pattern - text_chunks = re.findall(self.compiled_pattern, text) - # all chunks of text are encoded separately, then results are joined - ids = [] - for chunk in text_chunks: - chunk_bytes = chunk.encode("utf-8") # raw bytes - chunk_ids = self._encode_chunk(chunk_bytes) - ids.extend(chunk_ids) - return ids - -# ----------------------------------------------------------------------------- -# Faster Python tokenizer, optimized version of the reference tokenizer - -def fast_merge_inplace(ids, pair, idx): - """ - In the list of integers (ids), replace all consecutive occurrences - of pair with the new integer token idx in place - Example: ids=[1, 2, 3, 1, 2], pair=(1, 2), idx=4 -> [4, 3, 4] - """ - # Find all positions where the pair occurs - i = 0 - while i < len(ids) - 1: - if ids[i] == pair[0] and ids[i+1] == pair[1]: - ids[i] = idx - ids.pop(i+1) - else: - i += 1 - return ids - - -class FastRegexTokenizer: - - def __init__(self, pattern=None): - """ - - pattern: optional string to override the default (GPT-4 split pattern) - - special_tokens: str -> int dictionary of special tokens - example: {'<|endoftext|>': 100257} - """ - self.pattern = GPT4_SPLIT_PATTERN if pattern is None else pattern - self.compiled_pattern = re.compile(self.pattern) - self.special_tokens = {} - self.inverse_special_tokens = {} - self.merges = {} - self.vocab = self._build_vocab() - - def _build_vocab(self): - # vocab is simply and deterministically derived from merges - vocab = {idx: bytes([idx]) for idx in range(256)} - for (p0, p1), idx in self.merges.items(): - vocab[idx] = vocab[p0] + vocab[p1] - for special, idx in self.special_tokens.items(): - vocab[idx] = special.encode("utf-8") - return vocab - - def train(self, text, vocab_size, verbose=False): - """ - A number of optimizations are introduced: - - delete function call overhead by inlining functions - - modifying list of ids in place with .pop() instead of creating a new list - - collapse identical chunks to just the unique ones - - update counts more cleverly - only around the affected chunks - """ - assert vocab_size >= 256 - num_merges = vocab_size - 256 - - # split the text up into text chunks - text_chunks = re.findall(self.compiled_pattern, text) - - # many, many chunks are identical, so we can "collapse" them to just the unique ones - counts = Counter(text_chunks) - unique_chunks = [ch for ch, count in counts.items()] - chunk_counts = [count for ch, count in counts.items()] - - # input text preprocessing - ids = [list(ch.encode("utf-8")) for ch in unique_chunks] - # iteratively merge the most common pairs to create new tokens - merges = {} # (int, int) -> int - vocab = {idx: bytes([idx]) for idx in range(256)} # idx -> bytes - - # Initial count: build stats and position tracking - stats = defaultdict(int) - positions = defaultdict(set) # pair -> set of chunk indices that contain this pair - - for chunk_idx, (chunk_ids, count) in enumerate(zip(ids, chunk_counts)): - for pair in zip(chunk_ids, chunk_ids[1:]): - stats[pair] += count - positions[pair].add(chunk_idx) - - for i in range(num_merges): - if not stats: - break - - # find the pair with the highest count - pair = max(stats, key=stats.get) - # mint a new token: assign it the next available id - idx = 256 + i - - # Get chunks that contain this pair - affected_chunks = positions[pair] - - # Track count changes for incremental update - count_changes = defaultdict(int) - - # Replace all occurrences of pair in affected chunks only - for chunk_idx in affected_chunks: - chunk_ids = ids[chunk_idx] - chunk_count = chunk_counts[chunk_idx] - ix = 0 - while ix < len(chunk_ids) - 1: - if chunk_ids[ix] == pair[0] and chunk_ids[ix+1] == pair[1]: - # Track what pairs are being removed/added - # Remove: (prev, A), (A, B), (B, next) - if ix > 0: - old_left = (chunk_ids[ix-1], chunk_ids[ix]) - count_changes[old_left] -= chunk_count - - # The merged pair disappears - count_changes[pair] -= chunk_count - - if ix + 2 < len(chunk_ids): - old_right = (chunk_ids[ix+1], chunk_ids[ix+2]) - count_changes[old_right] -= chunk_count - - # Apply the merge - chunk_ids[ix] = idx - chunk_ids.pop(ix+1) - - # Add: (prev, C), (C, next) - if ix > 0: - new_left = (chunk_ids[ix-1], chunk_ids[ix]) - count_changes[new_left] += chunk_count - - if ix + 1 < len(chunk_ids): - new_right = (chunk_ids[ix], chunk_ids[ix+1]) - count_changes[new_right] += chunk_count - else: - ix += 1 - - # Apply incremental changes to stats and positions - for changed_pair, delta in count_changes.items(): - if changed_pair == pair: - # The merged pair should disappear completely - continue - - stats[changed_pair] += delta - - # Update positions for changed pairs - only check affected chunks - for chunk_idx in affected_chunks: - chunk_ids = ids[chunk_idx] - contains_pair = any((chunk_ids[j], chunk_ids[j+1]) == changed_pair - for j in range(len(chunk_ids) - 1)) - if contains_pair: - positions[changed_pair].add(chunk_idx) - else: - positions[changed_pair].discard(chunk_idx) - - # Remove the merged pair completely - del stats[pair] - del positions[pair] - - # save the merge - merges[pair] = idx - vocab[idx] = vocab[pair[0]] + vocab[pair[1]] - - # save class variables - self.merges = merges # used in encode() - self.vocab = vocab # used in decode() - - def register_special_tokens(self, special_tokens): - # special_tokens is a dictionary of str -> int - # example: {"<|endoftext|>": 100257} - self.special_tokens = special_tokens - self.inverse_special_tokens = {v: k for k, v in special_tokens.items()} - - def decode(self, ids): - # given ids (list of integers), return Python string - part_bytes = [] - for idx in ids: - if idx in self.vocab: - part_bytes.append(self.vocab[idx]) - elif idx in self.inverse_special_tokens: - part_bytes.append(self.inverse_special_tokens[idx].encode("utf-8")) - else: - raise ValueError(f"invalid token id: {idx}") - text_bytes = b"".join(part_bytes) - text = text_bytes.decode("utf-8", errors="replace") - return text - - def _encode_chunk(self, text_bytes): - # return the token ids - # let's begin. first, convert all bytes to integers in range 0..255 - ids = list(text_bytes) - while len(ids) >= 2: - # find the pair with the lowest merge index - stats = get_stats(ids) - pair = min(stats, key=lambda p: self.merges.get(p, float("inf"))) - # subtle: if there are no more merges available, the key will - # result in an inf for every single pair, and the min will be - # just the first pair in the list, arbitrarily - # we can detect this terminating case by a membership check - if pair not in self.merges: - break # nothing else can be merged anymore - # otherwise let's merge the best pair (lowest merge index) - idx = self.merges[pair] - ids = fast_merge_inplace(ids, pair, idx) - return ids - - def encode_ordinary(self, text): - """Encoding that ignores any special tokens.""" - # split text into chunks of text by categories defined in regex pattern - text_chunks = re.findall(self.compiled_pattern, text) - # all chunks of text are encoded separately, then results are joined - ids = [] - for chunk in text_chunks: - chunk_bytes = chunk.encode("utf-8") # raw bytes - chunk_ids = self._encode_chunk(chunk_bytes) - ids.extend(chunk_ids) - return ids - -# ----------------------------------------------------------------------------- -# HuggingFace tokenizer -from tokenizers import Tokenizer as HFTokenizer -from tokenizers import pre_tokenizers, decoders, Regex -from tokenizers.models import BPE -from tokenizers.trainers import BpeTrainer - -class HuggingFaceTokenizer: - """Light wrapper around HuggingFace Tokenizer for some utilities""" - - def __init__(self, tokenizer): - self.tokenizer = tokenizer - - @classmethod - def train_from_iterator(cls, text_iterator, vocab_size): - # train from an iterator of text - # Configure the HuggingFace Tokenizer - tokenizer = HFTokenizer(BPE( - byte_fallback=True, # needed! - unk_token=None, - fuse_unk=False, - )) - # Normalizer: None - tokenizer.normalizer = None - # Pre-tokenizer: GPT-4 style - gpt4_split_regex = Regex(GPT4_SPLIT_PATTERN) # huggingface demands that you wrap it in Regex!! - tokenizer.pre_tokenizer = pre_tokenizers.Sequence([ - pre_tokenizers.Split(pattern=gpt4_split_regex, behavior="isolated", invert=False), - pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False) - ]) - # Decoder: ByteLevel (it pairs together with the ByteLevel pre-tokenizer) - tokenizer.decoder = decoders.ByteLevel() - # Post-processor: None - tokenizer.post_processor = None - # Trainer: BPE - trainer = BpeTrainer( - vocab_size=vocab_size, - show_progress=True, - min_frequency=0, # no minimum frequency - initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), - special_tokens=[], # no special tokens - ) - # Kick off the training - tokenizer.train_from_iterator(text_iterator, trainer) - return cls(tokenizer) - - def encode_ordinary(self, text): - ids = self.tokenizer.encode(text, add_special_tokens=False).ids - return ids - -# ----------------------------------------------------------------------------- -# Test all of the above - -@pytest.fixture(scope="module") -def enwik8_path(): - """Fixture to download and cache enwik8 dataset.""" - import os - import zipfile - from nanochat.common import get_base_dir - base_dir = get_base_dir() - # download and unzip enwik8 to .cache directory - enwik8_url = "https://mattmahoney.net/dc/enwik8.zip" - enwik8_local_path = os.path.join(base_dir, "enwik8") - enwik8_local_path_zip = os.path.join(base_dir, "enwik8.zip") - if not os.path.exists(enwik8_local_path): - print(f"Downloading enwik8 to {enwik8_local_path_zip}") - import requests - response = requests.get(enwik8_url) - with open(enwik8_local_path_zip, "wb") as f: - f.write(response.content) - with zipfile.ZipFile(enwik8_local_path_zip, "r") as zip_ref: - zip_ref.extractall(base_dir) - print(f"Unzipped enwik8 to {enwik8_local_path}") - os.remove(enwik8_local_path_zip) - print(f"Removed {enwik8_local_path_zip}") - else: - print(f"Using existing enwik8 at {enwik8_local_path}") - return enwik8_local_path - - -@pytest.fixture(scope="module") -def enwik8_small(enwik8_path): - """Fixture providing 100KB of enwik8 for quick tests.""" - with open(enwik8_path, "r", encoding="utf-8") as f: - return f.read(100_000) - -@pytest.fixture(scope="module") -def enwik8_large(enwik8_path): - """Fixture providing 10MB of enwik8 for performance tests.""" - with open(enwik8_path, "r", encoding="utf-8") as f: - return f.read(10**7) - -def time_function(func, *args, **kwargs): - """Time a function call and return the result and elapsed time""" - start_time = time.time() - result = func(*args, **kwargs) - end_time = time.time() - elapsed = end_time - start_time - return result, elapsed - -def test_correctness(enwik8_small): - """Test that all tokenizer implementations produce the same results.""" - text = enwik8_small - encode_text = text - vocab_size = 256 + 20 # 20 merges - - # Train slow reference - print("\nTraining slow reference...") - slow_reference_tokenizer = RegexTokenizer() - ambiguous_flag, slow_reference_train_time = time_function(slow_reference_tokenizer.train, text, vocab_size) - slow_reference_ids, slow_reference_encode_time = time_function(slow_reference_tokenizer.encode_ordinary, encode_text) - print(f"Slow reference train time: {slow_reference_train_time:.4f}s") - print(f"Slow reference encode time: {slow_reference_encode_time:.4f}s") - print(slow_reference_ids[:20]) - - if ambiguous_flag: - print("‼️ WARNING: merge order was detected to be ambiguous given current text and vocab size") - print("The implementation could be correct but we might see different results below") - else: - print("✅ Merge order is NOT ambiguous") - - # Train fast reference - print("\nTraining fast reference...") - fast_reference_tokenizer = FastRegexTokenizer() - _, fast_reference_train_time = time_function(fast_reference_tokenizer.train, text, vocab_size) - fast_reference_ids, fast_reference_encode_time = time_function(fast_reference_tokenizer.encode_ordinary, encode_text) - print(f"Fast reference train time: {fast_reference_train_time:.4f}s") - print(f"Fast reference encode time: {fast_reference_encode_time:.4f}s") - print(fast_reference_ids[:20]) - - # Assert fast equals slow - assert fast_reference_ids == slow_reference_ids, "Fast reference should match slow reference" - print("✅ Fast == Slow") - - # Train HuggingFace - print("\nTraining HuggingFace...") - hf_tokenizer, hf_train_time = time_function(HuggingFaceTokenizer.train_from_iterator, [text], vocab_size) - hf_ids, hf_encode_time = time_function(hf_tokenizer.encode_ordinary, encode_text) - print(f"HuggingFace train time: {hf_train_time:.4f}s") - print(f"HuggingFace encode time: {hf_encode_time:.4f}s") - print(hf_ids[:20]) - - # HuggingFace has a different byte order, so we need custom matching - def custom_match(ids1, ids2): - perm = {} - for x, y in zip(ids1, ids2): - if x < 256: - if x in perm: - if perm[x] != y: - return False - perm[x] = y - if x >= 256 and x != y: - return False - return True - - assert custom_match(hf_ids, fast_reference_ids), "HuggingFace should match fast reference" - print("✅ HuggingFace == Fast") - - # Finally use our own Rust implementation - print("\nTraining rustbpe...") - rustbpe_tokenizer = rustbpe.Tokenizer() - _, rustbpe_train_time = time_function(rustbpe_tokenizer.train_from_iterator, [text], vocab_size) - rustbpe_ids, rustbpe_encode_time = time_function(rustbpe_tokenizer.encode, encode_text) - print(f"RustBPE train time: {rustbpe_train_time:.4f}s") - print(f"RustBPE encode time: {rustbpe_encode_time:.4f}s") - print(rustbpe_ids[:20]) - - assert rustbpe_ids == fast_reference_ids, "RustBPE should match fast reference" - print("✅ RustBPE == Fast") - - # Now export rustbpe to tiktoken for more efficient inference - print("\nTesting tiktoken export...") - pattern = rustbpe_tokenizer.get_pattern() - mergeable_ranks_list = rustbpe_tokenizer.get_mergeable_ranks() - mergeable_ranks = {bytes(k): v for k, v in mergeable_ranks_list} - enc = tiktoken.Encoding( - name="rustbpe", - pat_str=pattern, - mergeable_ranks=mergeable_ranks, - special_tokens={}, - ) - tiktoken_ids, tiktoken_encode_time = time_function(enc.encode, encode_text) - print(f"Tiktoken encode time: {tiktoken_encode_time:.4f}s") - print(tiktoken_ids[:20]) - - assert tiktoken_ids == rustbpe_ids, "Tiktoken should match RustBPE" - print("✅ Tiktoken == RustBPE") - - -@pytest.mark.slow -def test_training_performance(enwik8_large): - """Use a bigger dataset and compare the training speed of the optimized tokenizers (Python, Rust, HuggingFace).""" - text = enwik8_large - vocab_size = 2048 - print(f"\nText length: {len(text)}") - - # Commenting out because it's just way too slow to matter - # Train optimized python version - # print("Training optimized python version...") - # optimized_python_tokenizer = FastRegexTokenizer() - # _, optimized_python_train_time = time_function(optimized_python_tokenizer.train, text, vocab_size) - # print(f"Optimized python train time: {optimized_python_train_time:.4f}s") - - # Train rustbpe - print("\nTraining rustbpe...") - rustbpe_tokenizer = rustbpe.Tokenizer() - _, rustbpe_train_time = time_function(rustbpe_tokenizer.train_from_iterator, [text], vocab_size) - print(f"RustBPE train time: {rustbpe_train_time:.4f}s") - assert rustbpe_train_time > 0, "Training should take some time" - - # Train HuggingFace - print("\nTraining HuggingFace...") - hf_tokenizer, hf_train_time = time_function(HuggingFaceTokenizer.train_from_iterator, [text], vocab_size) - print(f"HuggingFace train time: {hf_train_time:.4f}s") - assert hf_train_time > 0, "Training should take some time" - - # Print comparison - print(f"\n📊 Performance comparison:") - print(f" RustBPE: {rustbpe_train_time:.4f}s") - print(f" HuggingFace: {hf_train_time:.4f}s") - print(f" Speedup: {hf_train_time/rustbpe_train_time:.2f}x") - -def test_interface(enwik8_small): - """Test the RustBPETokenizer interface for training, encoding, decoding, and serialization.""" - import tempfile - from nanochat.tokenizer import RustBPETokenizer - - # Simple train test - vocab_size = 300 - tok = RustBPETokenizer.train_from_iterator([enwik8_small], vocab_size) - assert tok.get_vocab_size() == vocab_size, f"Expected vocab size {vocab_size}, got {tok.get_vocab_size()}" - print(f"✅ Trained tokenizer with vocab size {vocab_size}") - - # Encode/decode text - encode_text = "Hello world! How are you? 🙃" - ids = tok.encode(encode_text) - print(f"\nInput text: {encode_text}") - print(f"IDs: {ids}") - decoded = tok.decode(ids) - print(f"Decoded: {decoded}") - assert decoded == encode_text, f"Decoded text doesn't match: {decoded} != {encode_text}" - print("✅ Encode/decode test passed") - - # Encode batch test - ids_new = tok.encode([encode_text, encode_text]) - assert all(x == ids for x in ids_new), "Batch encoding should produce identical results" - print("✅ Encode batch OK") - - # append/prepend functionality - ids_special = tok.encode(encode_text, prepend="<|bos|>", append="<|bos|>") - bos_token_id = tok.encode_special("<|bos|>") - assert ids_special == [bos_token_id] + ids + [bos_token_id], "Special tokens not correctly added" - print("✅ append/prepend OK") - - # Save/load test through a temporary directory - with tempfile.TemporaryDirectory() as tmp_dir: - tok.save(tmp_dir) - tok_reloaded = RustBPETokenizer.from_directory(tmp_dir) - ids_reloaded = tok_reloaded.encode(encode_text) - assert ids_reloaded == ids, "Reloaded tokenizer should produce same results" - print("✅ Save/load through temporary directory OK") diff --git a/uv.lock b/uv.lock index 4e9b0bdc..4e02a6c9 100644 --- a/uv.lock +++ b/uv.lock @@ -341,18 +341,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/42/14/42b2651a2f46b022ccd948bca9f2d5af0fd8929c4eec235b8d6d844fbe67/filelock-3.19.1-py3-none-any.whl", hash = "sha256:d38e30481def20772f5baf097c122c3babc4fcdb7e14e57049eb9d88c6dc017d", size = 15988, upload-time = "2025-08-14T16:56:01.633Z" }, ] -[[package]] -name = "files-to-prompt" -version = "0.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b9/4f/81fc86a88dc9e0cf6ea1ac2c561c0ac48b46d314cbbc2db5c8844b4b448b/files_to_prompt-0.6.tar.gz", hash = "sha256:9af57eecbdb29d3cce034c186493ffc6c1205ea4f5abde6fb32ccb1d96eae40c", size = 12236, upload-time = "2025-02-19T05:58:28.2Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/99/0efff50ce810119d99eaa2fc0c7bbf66e4197e2defb89242f6e848004902/files_to_prompt-0.6-py3-none-any.whl", hash = "sha256:83d9a8b33246a10233218716a5c78034da4f5614748eda2f0ab94f1117801337", size = 10873, upload-time = "2025-02-19T05:58:26.728Z" }, -] - [[package]] name = "frozenlist" version = "1.7.0" @@ -616,30 +604,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4f/65/6079a46068dfceaeabb5dcad6d674f5f5c61a6fa5673746f42a9f4c233b3/MarkupSafe-3.0.2-cp313-cp313t-win_amd64.whl", hash = "sha256:e444a31f8db13eb18ada366ab3cf45fd4b31e4db1236a4448f68778c1d1a5a2f", size = 15739, upload-time = "2024-10-18T15:21:42.784Z" }, ] -[[package]] -name = "maturin" -version = "1.9.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "tomli", marker = "python_full_version < '3.11' or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/13/7c/b11b870fc4fd84de2099906314ce45488ae17be32ff5493519a6cddc518a/maturin-1.9.4.tar.gz", hash = "sha256:235163a0c99bc6f380fb8786c04fd14dcf6cd622ff295ea3de525015e6ac40cf", size = 213647, upload-time = "2025-08-27T11:37:57.079Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/90/0d99389eea1939116fca841cad0763600c8d3183a02a9478d066736c60e8/maturin-1.9.4-py3-none-linux_armv6l.whl", hash = "sha256:6ff37578e3f5fdbe685110d45f60af1f5a7dfce70a1e26dfe3810af66853ecae", size = 8276133, upload-time = "2025-08-27T11:37:23.325Z" }, - { url = "https://files.pythonhosted.org/packages/f4/ed/c8ec68b383e50f084bf1fa9605e62a90cd32a3f75d9894ed3a6e5d4cc5b3/maturin-1.9.4-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:f3837bb53611b2dafa1c090436c330f2d743ba305ef00d8801a371f4495e7e1b", size = 15994496, upload-time = "2025-08-27T11:37:27.092Z" }, - { url = "https://files.pythonhosted.org/packages/84/4e/401ff5f3cfc6b123364d4b94379bf910d7baee32c9c95b72784ff2329357/maturin-1.9.4-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:4227d627d8e3bfe45877a8d65e9d8351a9d01434549f0da75d2c06a1b570de58", size = 8362228, upload-time = "2025-08-27T11:37:31.181Z" }, - { url = "https://files.pythonhosted.org/packages/51/8e/c56176dd360da9650c62b8a5ecfb85432cf011e97e46c186901e6996002e/maturin-1.9.4-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:1bb2aa0fa29032e9c5aac03ac400396ddea12cadef242f8967e9c8ef715313a1", size = 8271397, upload-time = "2025-08-27T11:37:33.672Z" }, - { url = "https://files.pythonhosted.org/packages/d2/46/001fcc5c6ad509874896418d6169a61acd619df5b724f99766308c44a99f/maturin-1.9.4-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:a0868d52934c8a5d1411b42367633fdb5cd5515bec47a534192282167448ec30", size = 8775625, upload-time = "2025-08-27T11:37:35.86Z" }, - { url = "https://files.pythonhosted.org/packages/b4/2e/26fa7574f01c19b7a74680fd70e5bae2e8c40fed9683d1752e765062cc2b/maturin-1.9.4-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:68b7b833b25741c0f553b78e8b9e095b31ae7c6611533b3c7b71f84c2cb8fc44", size = 8051117, upload-time = "2025-08-27T11:37:38.278Z" }, - { url = "https://files.pythonhosted.org/packages/73/ee/ca7308832d4f5b521c1aa176d9265f6f93e0bd1ad82a90fd9cd799f6b28c/maturin-1.9.4-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:08dc86312afee55af778af919818632e35d8d0464ccd79cb86700d9ea560ccd7", size = 8132122, upload-time = "2025-08-27T11:37:40.499Z" }, - { url = "https://files.pythonhosted.org/packages/45/e8/c623955da75e801a06942edf1fdc4e772a9e8fbc1ceebbdc85d59584dc10/maturin-1.9.4-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:ef20ffdd943078c4c3699c29fb2ed722bb6b4419efdade6642d1dbf248f94a70", size = 10586762, upload-time = "2025-08-27T11:37:42.718Z" }, - { url = "https://files.pythonhosted.org/packages/3c/4b/19ad558fdf54e151b1b4916ed45f1952ada96684ee6db64f9cd91cabec09/maturin-1.9.4-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:368e958468431dfeec80f75eea9639b4356d8c42428b0128444424b083fecfb0", size = 8926988, upload-time = "2025-08-27T11:37:45.492Z" }, - { url = "https://files.pythonhosted.org/packages/7e/27/153ad15eccae26921e8a01812da9f3b7f9013368f8f92c36853f2043b2a3/maturin-1.9.4-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:273f879214f63f79bfe851cd7d541f8150bdbfae5dfdc3c0c4d125d02d1f41b4", size = 8536758, upload-time = "2025-08-27T11:37:48.213Z" }, - { url = "https://files.pythonhosted.org/packages/43/e3/f304c3bdc3fba9adebe5348d4d2dd015f1152c0a9027aaf52cae0bb182c8/maturin-1.9.4-py3-none-win32.whl", hash = "sha256:ed2e54d132ace7e61829bd49709331007dd9a2cc78937f598aa76a4f69b6804d", size = 7265200, upload-time = "2025-08-27T11:37:50.881Z" }, - { url = "https://files.pythonhosted.org/packages/14/14/f86d0124bf1816b99005c058a1dbdca7cb5850d9cf4b09dcae07a1bc6201/maturin-1.9.4-py3-none-win_amd64.whl", hash = "sha256:8e450bb2c9afdf38a0059ee2e1ec2b17323f152b59c16f33eb9c74edaf1f9f79", size = 8237391, upload-time = "2025-08-27T11:37:53.23Z" }, - { url = "https://files.pythonhosted.org/packages/3f/25/8320fc2591e45b750c3ae71fa596b47aefa802d07d6abaaa719034a85160/maturin-1.9.4-py3-none-win_arm64.whl", hash = "sha256:7a6f980a9b67a5c13c844c268eabd855b54a6a765df4b4bb07d15a990572a4c9", size = 6988277, upload-time = "2025-08-27T11:37:55.429Z" }, -] - [[package]] name = "mpmath" version = "1.3.0" @@ -772,36 +736,38 @@ wheels = [ [[package]] name = "nanochat" version = "0.1.0" -source = { editable = "." } +source = { virtual = "." } dependencies = [ { name = "datasets" }, { name = "fastapi" }, - { name = "files-to-prompt" }, { name = "psutil" }, + { name = "python-dotenv" }, { name = "regex" }, + { name = "rustbpe" }, { name = "setuptools" }, { name = "tiktoken" }, { name = "tokenizers" }, - { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "extra == 'extra-8-nanochat-gpu'" }, - { name = "torch", version = "2.9.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu')" }, - { name = "torch", version = "2.9.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.0", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.1", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.1+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.1+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "extra == 'extra-8-nanochat-gpu'" }, + { name = "transformers" }, { name = "uvicorn" }, { name = "wandb" }, ] [package.optional-dependencies] cpu = [ - { name = "torch", version = "2.9.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "torch", version = "2.9.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.1", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "torch", version = "2.9.1+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] gpu = [ - { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" } }, + { name = "torch", version = "2.9.1+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" } }, ] [package.dev-dependencies] dev = [ - { name = "maturin" }, { name = "pytest" }, ] @@ -809,25 +775,24 @@ dev = [ requires-dist = [ { name = "datasets", specifier = ">=4.0.0" }, { name = "fastapi", specifier = ">=0.117.1" }, - { name = "files-to-prompt", specifier = ">=0.6" }, { name = "psutil", specifier = ">=7.1.0" }, + { name = "python-dotenv", specifier = ">=1.2.1" }, { name = "regex", specifier = ">=2025.9.1" }, + { name = "rustbpe", specifier = ">=0.1.0" }, { name = "setuptools", specifier = ">=80.9.0" }, { name = "tiktoken", specifier = ">=0.11.0" }, { name = "tokenizers", specifier = ">=0.22.0" }, { name = "torch", specifier = ">=2.8.0" }, - { name = "torch", marker = "extra == 'cpu'", specifier = ">=2.8.0", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "nanochat", extra = "cpu" } }, - { name = "torch", marker = "extra == 'gpu'", specifier = ">=2.8.0", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "nanochat", extra = "gpu" } }, + { name = "torch", marker = "extra == 'cpu'", specifier = ">=2.9.1", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "nanochat", extra = "cpu" } }, + { name = "torch", marker = "extra == 'gpu'", specifier = ">=2.9.1", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "nanochat", extra = "gpu" } }, + { name = "transformers", specifier = ">=4.57.3" }, { name = "uvicorn", specifier = ">=0.36.0" }, { name = "wandb", specifier = ">=0.21.3" }, ] provides-extras = ["cpu", "gpu"] [package.metadata.requires-dev] -dev = [ - { name = "maturin", specifier = ">=1.9.4" }, - { name = "pytest", specifier = ">=8.0.0" }, -] +dev = [{ name = "pytest", specifier = ">=8.0.0" }] [[package]] name = "networkx" @@ -949,7 +914,7 @@ name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cublas-cu12", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/fa/41/e79269ce215c857c935fd86bcfe91a451a584dfc27f1e068f568b9ad1ab7/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8", size = 705026878, upload-time = "2025-06-06T21:52:51.348Z" }, @@ -962,7 +927,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211, upload-time = "2025-03-07T01:44:56.873Z" }, @@ -994,9 +959,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cusparse-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cublas-cu12", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusparse-cu12", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841, upload-time = "2025-03-07T01:46:54.356Z" }, @@ -1009,7 +974,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129, upload-time = "2025-03-07T01:47:40.407Z" }, @@ -1029,11 +994,11 @@ wheels = [ [[package]] name = "nvidia-nccl-cu12" -version = "2.27.3" +version = "2.27.5" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/4b/7b/8354b784cf73b0ba51e566b4baba3ddd44fe8288a3d39ef1e06cd5417226/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9ddf1a245abc36c550870f26d537a9b6087fb2e2e3d6e0ef03374c6fd19d984f", size = 322397768, upload-time = "2025-06-03T21:57:30.234Z" }, - { url = "https://files.pythonhosted.org/packages/5c/5b/4e4fff7bad39adf89f735f2bc87248c81db71205b62bcc0d5ca5b606b3c3/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039", size = 322364134, upload-time = "2025-06-03T21:58:04.013Z" }, + { url = "https://files.pythonhosted.org/packages/bb/1c/857979db0ef194ca5e21478a0612bcdbbe59458d7694361882279947b349/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:31432ad4d1fb1004eb0c56203dc9bc2178a1ba69d1d9e02d64a6938ab5e40e7a", size = 322400625, upload-time = "2025-06-26T04:11:04.496Z" }, + { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, ] [[package]] @@ -1046,6 +1011,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ed/d7/34f02dad2e30c31b10a51f6b04e025e5dd60e5f936af9045a9b858a05383/nvidia_nvjitlink_cu12-12.8.93-py3-none-win_amd64.whl", hash = "sha256:bd93fbeeee850917903583587f4fc3a4eafa022e34572251368238ab5e6bd67f", size = 268553710, upload-time = "2025-03-07T01:56:24.13Z" }, ] +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.3.20" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/9d/3dd98852568fb845ec1f7902c90a22b240fe1cbabda411ccedf2fd737b7b/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b0b960da3842212758e4fa4696b94f129090b30e5122fea3c5345916545cff0", size = 124484616, upload-time = "2025-08-04T20:24:59.172Z" }, + { url = "https://files.pythonhosted.org/packages/3b/6c/99acb2f9eb85c29fc6f3a7ac4dccfd992e22666dd08a642b303311326a97/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d00f26d3f9b2e3c3065be895e3059d6479ea5c638a3f38c9fec49b1b9dd7c1e5", size = 124657145, upload-time = "2025-08-04T20:25:19.995Z" }, +] + [[package]] name = "nvidia-nvtx-cu12" version = "12.8.90" @@ -1434,6 +1408,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f0/26/19cadc79a718c5edbec86fd4919a6b6d3f681039a2f6d66d14be94e75fb9/python_dotenv-1.2.1.tar.gz", hash = "sha256:42667e897e16ab0d66954af0e60a9caa94f0fd4ecf3aaf6d2d260eec1aa36ad6", size = 44221, upload-time = "2025-10-26T15:12:10.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/1b/a298b06749107c305e1fe0f814c6c74aea7b2f1e10989cb30f544a1b3253/python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61", size = 21230, upload-time = "2025-10-26T15:12:09.109Z" }, +] + [[package]] name = "pytz" version = "2025.2" @@ -1581,6 +1564,69 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" }, ] +[[package]] +name = "rustbpe" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/03/2e/f16e179ad1e185f0bb5a8fc2376fff05d1eeefcb6d8a77ee04306e8a42ae/rustbpe-0.1.0.tar.gz", hash = "sha256:18765f62ac579a9ff9e89c611f9c9b9e46bd1adde9be3f59c00b6eb4e1f28b3a", size = 29723, upload-time = "2026-01-03T22:24:11.872Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/47/d4/cdb13041ebfc5e10b98fa0de1d631bbce5476fe265c5e97516c344bbf44d/rustbpe-0.1.0-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:15f2e126a2b9fde264598f9317b1ec7f20ed3cf1e49cae081f7d6b7c1064655c", size = 1010629, upload-time = "2026-01-03T22:23:33.475Z" }, + { url = "https://files.pythonhosted.org/packages/3d/dd/401cb29be42b3efa80a8499ac41532bb1aa488c16cb1921dca880f0d75ce/rustbpe-0.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:34dc37f2771517813b70be47e26d5861ba69298f299dfd24409882e76bd1ccad", size = 951962, upload-time = "2026-01-03T22:23:35.025Z" }, + { url = "https://files.pythonhosted.org/packages/3b/4c/1f17ebab894dcb72e5b15389ad4b06b74637745163371339f32f23cdec76/rustbpe-0.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d4e7543962c932c3dcc29f5a02855ad613e9b05bc93c898fb99e4c7928598c7b", size = 1034660, upload-time = "2026-01-03T22:23:36.353Z" }, + { url = "https://files.pythonhosted.org/packages/ba/95/ab37d09e7b51b3ae49ef5af885c3d7c0244c72521f0740c6b55e1827a251/rustbpe-0.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7e90da31ce2708481c6b2169d9b8072917d7deefae4fb91cf25655dedac6afa", size = 1078818, upload-time = "2026-01-03T22:23:37.606Z" }, + { url = "https://files.pythonhosted.org/packages/94/65/79a4005477545c0661571e642d614177525e010ff9fc27cc10e579bde38d/rustbpe-0.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:9b92ff475c37a2b10391614c10ab2deb923bb029a8288e11e77618d5e5fcdc0e", size = 919699, upload-time = "2026-01-03T22:23:38.944Z" }, + { url = "https://files.pythonhosted.org/packages/16/c1/d4fadf70d1cc0914c812a9c7c1e5cce0813440f7d16082fdb399ec33748d/rustbpe-0.1.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:400be6ede8875d5ac0e0ac91dfba1ec7ea7d359353b0465da633576cf01c7de7", size = 1008245, upload-time = "2026-01-03T22:23:40.245Z" }, + { url = "https://files.pythonhosted.org/packages/8d/e1/ac7d4044dbee242bbcb7d9fc425f6ea8c52f984c7708cbb4cb9633976b96/rustbpe-0.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:dff3ffb6f05576a27732d2013f044ec6f137bc7bce6773a5e134cfc0c24dcc82", size = 949344, upload-time = "2026-01-03T22:23:41.664Z" }, + { url = "https://files.pythonhosted.org/packages/2a/7b/008e45858130eb803085d131a05e6e55c123a2b63b763ea08a45aa8b7673/rustbpe-0.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:92a0186ed815ccec376cca23c4bc5f209f6c67efeb101c1c935345cd63cc9eea", size = 1031915, upload-time = "2026-01-03T22:23:42.93Z" }, + { url = "https://files.pythonhosted.org/packages/1f/6e/d10c687670c42d34306713ae75d6477d6c32424bd251033bd9ff2a243ccd/rustbpe-0.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fec78edb30f3264d0db69ffd7ac333d695be76e4e672fd5301626787bc1220c2", size = 1076476, upload-time = "2026-01-03T22:23:43.899Z" }, + { url = "https://files.pythonhosted.org/packages/78/a8/f64b877d0a0239f4262a90d74ded014f1e2c4250c6273898280739177a7b/rustbpe-0.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:0f35a858c31faf09e6723fe2e8c020efcf4e036b7270ed151ca8538fad1fe0c5", size = 916888, upload-time = "2026-01-03T22:23:44.936Z" }, + { url = "https://files.pythonhosted.org/packages/a9/a3/7fe53c4dcd7d90a777424c61ac8072153ce47941066e0a247c020a4a663e/rustbpe-0.1.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:13e6aeaaf6e2f970ab577f32a6c49c8dd23517279253a37873ddc7f74fd30622", size = 1007207, upload-time = "2026-01-03T22:23:46.336Z" }, + { url = "https://files.pythonhosted.org/packages/a7/41/dee1474cfea594d7a9cebb42f683170f1f2d8af4473541c0a1f96dfaff76/rustbpe-0.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40763a0751ba8a595717f5015d18b0241e1af9930412e42d350380ba4601361b", size = 947913, upload-time = "2026-01-03T22:23:47.458Z" }, + { url = "https://files.pythonhosted.org/packages/a2/fd/c90bc3a3e823b8cafb85625ed37311987c20317168ea73d0ebaba54f8df2/rustbpe-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bfe8d24d0d71c16fb8ba5106e7d2be2c43211195a74ffa7e2c88cb98c07122e4", size = 1030968, upload-time = "2026-01-03T22:23:48.753Z" }, + { url = "https://files.pythonhosted.org/packages/fa/64/e15606774d2f13d1bdbdca4cd6e8fcd14fc0c3fb7ca7b00412c4ed0a8700/rustbpe-0.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:844f7f6c3bd59a9578b87ebc6bb60fe3ee47c8d8040a62488ce8e7eaeeb31319", size = 1075101, upload-time = "2026-01-03T22:23:50.041Z" }, + { url = "https://files.pythonhosted.org/packages/d7/26/8de98d90fd8765a1ea517b01897e05aa9932998e604bb9003e5e9b73be3c/rustbpe-0.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:b79b67d8db6a2fe3928918006569e73aee23e012b3b0b36fd4a2a85cc2c2161f", size = 914924, upload-time = "2026-01-03T22:23:51.31Z" }, + { url = "https://files.pythonhosted.org/packages/c6/63/a0475defd438cd6a4cd28b74ad8dd01bb7de6adafaa411968e758b0a9036/rustbpe-0.1.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:6a59c05a8123d3a8e8815106fd1938a9499d4fdaf5cf00351fa7d3b5cc4f8ad6", size = 1007322, upload-time = "2026-01-03T22:23:52.568Z" }, + { url = "https://files.pythonhosted.org/packages/81/72/18e762472a42d68820e2d1244655fd960e200e449136fabe3c32f6f2a1b1/rustbpe-0.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bc0a8dd8b30860e3a4889ab7cf7c04a2614f8fc77c191efde1500aa054484efa", size = 948256, upload-time = "2026-01-03T22:23:53.926Z" }, + { url = "https://files.pythonhosted.org/packages/16/07/3c0948db94fc454b62012ff8b3e74ad13f84bf8fbcfb84b402bfb786e82e/rustbpe-0.1.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3eeb8efae3d10a3b6640a1e2bdc7f1e55a15f867bdae9efb3d8f0757b01d9d3a", size = 1031258, upload-time = "2026-01-03T22:23:54.961Z" }, + { url = "https://files.pythonhosted.org/packages/fb/69/77355ca8baf0c5023994b3f11304822d07116567ea47893f90267c086f87/rustbpe-0.1.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2b0b77591c9e836df41ad0b30be9ec519a708c477cbf82eedaf839e7a9b10101", size = 1075321, upload-time = "2026-01-03T22:23:55.995Z" }, + { url = "https://files.pythonhosted.org/packages/5c/fe/5c529d92988be7df251de718a633054ecca2d5986a17759a6546a9f45c26/rustbpe-0.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:cc5bebc9990071e400bbe304304af3d5757522bb2a1177e2c3517f11ad28f0eb", size = 915136, upload-time = "2026-01-03T22:23:57.56Z" }, + { url = "https://files.pythonhosted.org/packages/af/d7/8f7215233acd67402f8bdf972daa3fbe9184b176348530b84ac40751a806/rustbpe-0.1.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b1229e70c2d091faf8c0a50951e2e734b3b810d1d2b7677cd49d86dc3853c283", size = 1031277, upload-time = "2026-01-03T22:23:58.55Z" }, + { url = "https://files.pythonhosted.org/packages/ff/1a/0b34c02138f28a984bc44fdc0dc10afc9137814b2a56b8cd4e5ae25b8601/rustbpe-0.1.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:02d123e72fe9253c92904bfe2ba35afc816576b2cdbb432a96001e75bafb888e", size = 1007777, upload-time = "2026-01-03T22:23:59.539Z" }, + { url = "https://files.pythonhosted.org/packages/bb/b1/da66ce14f43b23136c07183be03ddbc58654824455cce36c2bad38254aeb/rustbpe-0.1.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:a0df5172a813c982d31673d9de32dd053ddbb64ced2b97709a85d2e3c6a6cd28", size = 948400, upload-time = "2026-01-03T22:24:00.506Z" }, + { url = "https://files.pythonhosted.org/packages/05/d0/551dcfb8d314f4e0b60b86ab616bcaaf3a381f6e72f83f1211246528a7c1/rustbpe-0.1.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c525072e521a5cf729474a0ea6c83b1b16973b877098ee7060eac4bbacd46c7a", size = 1031325, upload-time = "2026-01-03T22:24:01.501Z" }, + { url = "https://files.pythonhosted.org/packages/4e/36/3f1730a6b8f4435b8cb2ceee2edb3be8357656e35f1f6549b5f387eb056a/rustbpe-0.1.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e0c9216f9e38558f0939f6e34cc78d5517d7a02026c1a35b271ca82e9b522539", size = 1075729, upload-time = "2026-01-03T22:24:02.528Z" }, + { url = "https://files.pythonhosted.org/packages/b4/03/aaa994e9a28cb7248c2cfc43a93c779ee7ac0e19cf9eae6717b63bbe6a8d/rustbpe-0.1.0-cp314-cp314-win_amd64.whl", hash = "sha256:b5ceb789bb93a82547c0ed7277ecc01047eaf0eeea6bbc0a21420e65e5fb553a", size = 915650, upload-time = "2026-01-03T22:24:03.71Z" }, + { url = "https://files.pythonhosted.org/packages/8c/68/3ab181ff8b12dcabdb256dffb82de0d8bf30c72ac3d188451ac5fa1cc643/rustbpe-0.1.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:88d1482ccadf5e29b524b13740b3a5e1f4e454a048684885d894fd1a9930617a", size = 1030995, upload-time = "2026-01-03T22:24:04.76Z" }, + { url = "https://files.pythonhosted.org/packages/96/a2/02498910b4852967fd4b6d77ce94542c5483f1551decb6911480229d116c/rustbpe-0.1.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ba8e08ed3cc7a7bf832f70c86c64a94d0112e8c526d55a1f40e53ede2ca14d22", size = 1031327, upload-time = "2026-01-03T22:24:09.246Z" }, + { url = "https://files.pythonhosted.org/packages/49/13/78d768a451dc9e634f933f2231b3fa9be524955ed84317b40e5528a2d906/rustbpe-0.1.0-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9f419fd428e8ffd2430945a694cb5177706550ee5c9b16737ba860ecccd5acff", size = 1075802, upload-time = "2026-01-03T22:24:10.573Z" }, +] + +[[package]] +name = "safetensors" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, + { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, + { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, + { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, + { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, + { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, + { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, + { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, + { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, + { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, + { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, + { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, + { url = "https://files.pythonhosted.org/packages/a7/6a/4d08d89a6fcbe905c5ae68b8b34f0791850882fc19782d0d02c65abbdf3b/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f4729811a6640d019a4b7ba8638ee2fd21fa5ca8c7e7bdf0fed62068fcaac737", size = 492430, upload-time = "2025-11-19T15:18:11.884Z" }, + { url = "https://files.pythonhosted.org/packages/dd/29/59ed8152b30f72c42d00d241e58eaca558ae9dbfa5695206e2e0f54c7063/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:12f49080303fa6bb424b362149a12949dfbbf1e06811a88f2307276b0c131afd", size = 503977, upload-time = "2025-11-19T15:18:17.523Z" }, + { url = "https://files.pythonhosted.org/packages/d3/0b/4811bfec67fa260e791369b16dab105e4bae82686120554cc484064e22b4/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0071bffba4150c2f46cae1432d31995d77acfd9f8db598b5d1a2ce67e8440ad2", size = 623890, upload-time = "2025-11-19T15:18:22.666Z" }, + { url = "https://files.pythonhosted.org/packages/58/5b/632a58724221ef03d78ab65062e82a1010e1bef8e8e0b9d7c6d7b8044841/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:473b32699f4200e69801bf5abf93f1a4ecd432a70984df164fc22ccf39c4a6f3", size = 531885, upload-time = "2025-11-19T15:18:27.146Z" }, +] + [[package]] name = "sentry-sdk" version = "2.35.2" @@ -1755,106 +1801,40 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6e/c2/61d3e0f47e2b74ef40a68b9e6ad5984f6241a942f7cd3bbfbdbd03861ea9/tomli-2.2.1-py3-none-any.whl", hash = "sha256:cb55c73c5f4408779d0cf3eef9f762b9c9f147a77de7b258bef0a5628adc85cc", size = 14257, upload-time = "2024-11-27T22:38:35.385Z" }, ] -[[package]] -name = "torch" -version = "2.8.0+cu128" -source = { registry = "https://download.pytorch.org/whl/cu128" } -resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform != 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform != 'linux'", - "python_full_version < '3.11' and sys_platform != 'linux'", -] -dependencies = [ - { name = "filelock", marker = "extra == 'extra-8-nanochat-gpu'" }, - { name = "fsspec", marker = "extra == 'extra-8-nanochat-gpu'" }, - { name = "jinja2", marker = "extra == 'extra-8-nanochat-gpu'" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cuda-cupti-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cuda-runtime-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cudnn-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cufft-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cufile-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-curand-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cusolver-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cusparse-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-cusparselt-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-nccl-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "nvidia-nvtx-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "setuptools", marker = "(python_full_version >= '3.12' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "sympy", marker = "extra == 'extra-8-nanochat-gpu'" }, - { name = "triton", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (platform_machine != 'x86_64' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'linux' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "typing-extensions", marker = "extra == 'extra-8-nanochat-gpu'" }, -] -wheels = [ - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:0c96999d15cf1f13dd7c913e0b21a9a355538e6cfc10861a17158320292f5954" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:43938e9a174c90e5eb9e906532b2f1e21532bbfa5a61b65193b4f54714d34f9e" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:039b9dcdd6bdbaa10a8a5cd6be22c4cb3e3589a341e5f904cbb571ca28f55bed" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp311-cp311-win_amd64.whl", hash = "sha256:34c55443aafd31046a7963b63d30bc3b628ee4a704f826796c865fdfd05bb596" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:4354fc05bb79b208d6995a04ca1ceef6a9547b1c4334435574353d381c55087c" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp312-cp312-win_amd64.whl", hash = "sha256:0ad925202387f4e7314302a1b4f8860fa824357f9b1466d7992bf276370ebcff" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:3a852369a38dec343d45ecd0bc3660f79b88a23e0c878d18707f7c13bf49538f" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp313-cp313-win_amd64.whl", hash = "sha256:9e20646802b7fc295c1f8b45fefcfc9fb2e4ec9cbe8593443cd2b9cc307c8405" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:4295a22d69408e93d25f51e8d5d579345b6b802383e9414b0f3853ed433d53ae" }, - { url = "https://download.pytorch.org/whl/cu128/torch-2.8.0%2Bcu128-cp313-cp313t-win_amd64.whl", hash = "sha256:970b4f4661fa7b44f6a7e6df65de7fc4a6fff2af610dc415c1d695ca5f1f37d2" }, -] - -[[package]] -name = "torch" -version = "2.9.0" -source = { registry = "https://download.pytorch.org/whl/cpu" } -resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'darwin'", - "python_full_version == '3.11.*' and sys_platform == 'darwin'", - "python_full_version < '3.11' and sys_platform == 'darwin'", -] -dependencies = [ - { name = "filelock", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "fsspec", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "jinja2", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (python_full_version >= '3.11' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (python_full_version < '3.11' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "setuptools", marker = "(python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (python_full_version < '3.12' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "sympy", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "typing-extensions", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, -] -wheels = [ - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:59484193b01299bf669520505a72b29d59a0028ae4c6d95f492938f186592208" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aa4483602586cc9a35d1cf33771a9977f05f642b9161518a289e36548a0b77c2" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:4de0ed8cbc457a506dbca40376e206a29efee10756a00f1f3404bf67ad737d04" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:259548471194ab63d7ea273873053a6e3cc23530c1510f01e9d7ad259187bbd0" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e24836d968b54ef4dfb05594001a61958711ac9224026291e4e3f92f83a6fd7f" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d8e2ab7f86010330bdcc39c8b2c795590cc75e37df4823cdaee2c98d6e3ff4a3" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a3e859039c985d8e3ea60d7a54ca7e97ea2ae15e31beced4f3260128a161bb01" }, -] - [[package]] name = "torch" version = "2.9.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform != 'linux'", "python_full_version == '3.11.*' and sys_platform == 'linux'", "python_full_version < '3.11' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform != 'linux'", - "python_full_version < '3.11' and sys_platform != 'linux'", ] dependencies = [ - { name = "filelock", marker = "(extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu')" }, - { name = "fsspec", marker = "(extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu')" }, - { name = "jinja2", marker = "(extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu')" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "setuptools", marker = "(python_full_version >= '3.12' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, - { name = "sympy", marker = "(extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu')" }, - { name = "typing-extensions", marker = "(extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu')" }, + { name = "filelock", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "fsspec", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "jinja2", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cublas-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cuda-cupti-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cuda-runtime-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cudnn-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cufft-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cufile-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-curand-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusolver-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusparse-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusparselt-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nccl-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvshmem-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvtx-cu12", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "sympy", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine == 'x86_64' and sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "typing-extensions", marker = "(sys_platform == 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/bb/86/245c240d2138c17ed572c943c289056c2721abab70810d772c6bf5495b28/torch-2.9.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:030bbfe367379ae6a4ae4042b6c44da25383343b8b3c68abaa9c7231efbaf2dd", size = 104213554, upload-time = "2025-10-15T15:45:59.798Z" }, @@ -1889,7 +1869,86 @@ wheels = [ [[package]] name = "torch" -version = "2.9.0+cpu" +version = "2.9.1" +source = { registry = "https://download.pytorch.org/whl/cpu" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform == 'darwin'", +] +dependencies = [ + { name = "filelock", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "fsspec", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "jinja2", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (python_full_version >= '3.11' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (python_full_version < '3.11' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (python_full_version < '3.12' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu') or (sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "sympy", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "typing-extensions", marker = "(sys_platform == 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, +] +wheels = [ + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp310-none-macosx_11_0_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp311-none-macosx_11_0_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp313-cp313t-macosx_11_0_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp313-none-macosx_11_0_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp314-cp314-macosx_11_0_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1-cp314-cp314t-macosx_11_0_arm64.whl" }, +] + +[[package]] +name = "torch" +version = "2.9.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform != 'linux'", + "python_full_version == '3.11.*' and sys_platform != 'linux'", + "python_full_version < '3.11' and sys_platform != 'linux'", +] +dependencies = [ + { name = "filelock", marker = "(sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "fsspec", marker = "(sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "jinja2", marker = "(sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "sympy", marker = "(sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "typing-extensions", marker = "(sys_platform != 'linux' and extra != 'extra-8-nanochat-cpu' and extra != 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/56/9577683b23072075ed2e40d725c52c2019d71a972fab8e083763da8e707e/torch-2.9.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:1cc208435f6c379f9b8fdfd5ceb5be1e3b72a6bdf1cb46c0d2812aa73472db9e", size = 104207681, upload-time = "2025-11-12T15:19:56.48Z" }, + { url = "https://files.pythonhosted.org/packages/38/45/be5a74f221df8f4b609b78ff79dc789b0cc9017624544ac4dd1c03973150/torch-2.9.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:9fd35c68b3679378c11f5eb73220fdcb4e6f4592295277fbb657d31fd053237c", size = 899794036, upload-time = "2025-11-12T15:21:01.886Z" }, + { url = "https://files.pythonhosted.org/packages/67/95/a581e8a382596b69385a44bab2733f1273d45c842f5d4a504c0edc3133b6/torch-2.9.1-cp310-cp310-win_amd64.whl", hash = "sha256:2af70e3be4a13becba4655d6cc07dcfec7ae844db6ac38d6c1dafeb245d17d65", size = 110969861, upload-time = "2025-11-12T15:21:30.145Z" }, + { url = "https://files.pythonhosted.org/packages/ad/51/1756dc128d2bf6ea4e0a915cb89ea5e730315ff33d60c1ff56fd626ba3eb/torch-2.9.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:a83b0e84cc375e3318a808d032510dde99d696a85fe9473fc8575612b63ae951", size = 74452222, upload-time = "2025-11-12T15:20:46.223Z" }, + { url = "https://files.pythonhosted.org/packages/15/db/c064112ac0089af3d2f7a2b5bfbabf4aa407a78b74f87889e524b91c5402/torch-2.9.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:62b3fd888277946918cba4478cf849303da5359f0fb4e3bfb86b0533ba2eaf8d", size = 104220430, upload-time = "2025-11-12T15:20:31.705Z" }, + { url = "https://files.pythonhosted.org/packages/56/be/76eaa36c9cd032d3b01b001e2c5a05943df75f26211f68fae79e62f87734/torch-2.9.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:d033ff0ac3f5400df862a51bdde9bad83561f3739ea0046e68f5401ebfa67c1b", size = 899821446, upload-time = "2025-11-12T15:20:15.544Z" }, + { url = "https://files.pythonhosted.org/packages/47/cc/7a2949e38dfe3244c4df21f0e1c27bce8aedd6c604a587dd44fc21017cb4/torch-2.9.1-cp311-cp311-win_amd64.whl", hash = "sha256:0d06b30a9207b7c3516a9e0102114024755a07045f0c1d2f2a56b1819ac06bcb", size = 110973074, upload-time = "2025-11-12T15:21:39.958Z" }, + { url = "https://files.pythonhosted.org/packages/1e/ce/7d251155a783fb2c1bb6837b2b7023c622a2070a0a72726ca1df47e7ea34/torch-2.9.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:52347912d868653e1528b47cafaf79b285b98be3f4f35d5955389b1b95224475", size = 74463887, upload-time = "2025-11-12T15:20:36.611Z" }, + { url = "https://files.pythonhosted.org/packages/0f/27/07c645c7673e73e53ded71705045d6cb5bae94c4b021b03aa8d03eee90ab/torch-2.9.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:da5f6f4d7f4940a173e5572791af238cb0b9e21b1aab592bd8b26da4c99f1cd6", size = 104126592, upload-time = "2025-11-12T15:20:41.62Z" }, + { url = "https://files.pythonhosted.org/packages/19/17/e377a460603132b00760511299fceba4102bd95db1a0ee788da21298ccff/torch-2.9.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:27331cd902fb4322252657f3902adf1c4f6acad9dcad81d8df3ae14c7c4f07c4", size = 899742281, upload-time = "2025-11-12T15:22:17.602Z" }, + { url = "https://files.pythonhosted.org/packages/b1/1a/64f5769025db846a82567fa5b7d21dba4558a7234ee631712ee4771c436c/torch-2.9.1-cp312-cp312-win_amd64.whl", hash = "sha256:81a285002d7b8cfd3fdf1b98aa8df138d41f1a8334fd9ea37511517cedf43083", size = 110940568, upload-time = "2025-11-12T15:21:18.689Z" }, + { url = "https://files.pythonhosted.org/packages/6e/ab/07739fd776618e5882661d04c43f5b5586323e2f6a2d7d84aac20d8f20bd/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:c0d25d1d8e531b8343bea0ed811d5d528958f1dcbd37e7245bc686273177ad7e", size = 74479191, upload-time = "2025-11-12T15:21:25.816Z" }, + { url = "https://files.pythonhosted.org/packages/20/60/8fc5e828d050bddfab469b3fe78e5ab9a7e53dda9c3bdc6a43d17ce99e63/torch-2.9.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:c29455d2b910b98738131990394da3e50eea8291dfeb4b12de71ecf1fdeb21cb", size = 104135743, upload-time = "2025-11-12T15:21:34.936Z" }, + { url = "https://files.pythonhosted.org/packages/f2/b7/6d3f80e6918213babddb2a37b46dbb14c15b14c5f473e347869a51f40e1f/torch-2.9.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:524de44cd13931208ba2c4bde9ec7741fd4ae6bfd06409a604fc32f6520c2bc9", size = 899749493, upload-time = "2025-11-12T15:24:36.356Z" }, + { url = "https://files.pythonhosted.org/packages/a6/47/c7843d69d6de8938c1cbb1eba426b1d48ddf375f101473d3e31a5fc52b74/torch-2.9.1-cp313-cp313-win_amd64.whl", hash = "sha256:545844cc16b3f91e08ce3b40e9c2d77012dd33a48d505aed34b7740ed627a1b2", size = 110944162, upload-time = "2025-11-12T15:21:53.151Z" }, + { url = "https://files.pythonhosted.org/packages/28/0e/2a37247957e72c12151b33a01e4df651d9d155dd74d8cfcbfad15a79b44a/torch-2.9.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5be4bf7496f1e3ffb1dd44b672adb1ac3f081f204c5ca81eba6442f5f634df8e", size = 74830751, upload-time = "2025-11-12T15:21:43.792Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f7/7a18745edcd7b9ca2381aa03353647bca8aace91683c4975f19ac233809d/torch-2.9.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:30a3e170a84894f3652434b56d59a64a2c11366b0ed5776fab33c2439396bf9a", size = 104142929, upload-time = "2025-11-12T15:21:48.319Z" }, + { url = "https://files.pythonhosted.org/packages/f4/dd/f1c0d879f2863ef209e18823a988dc7a1bf40470750e3ebe927efdb9407f/torch-2.9.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8301a7b431e51764629208d0edaa4f9e4c33e6df0f2f90b90e261d623df6a4e2", size = 899748978, upload-time = "2025-11-12T15:23:04.568Z" }, + { url = "https://files.pythonhosted.org/packages/1f/9f/6986b83a53b4d043e36f3f898b798ab51f7f20fdf1a9b01a2720f445043d/torch-2.9.1-cp313-cp313t-win_amd64.whl", hash = "sha256:2e1c42c0ae92bf803a4b2409fdfed85e30f9027a66887f5e7dcdbc014c7531db", size = 111176995, upload-time = "2025-11-12T15:22:01.618Z" }, + { url = "https://files.pythonhosted.org/packages/40/60/71c698b466dd01e65d0e9514b5405faae200c52a76901baf6906856f17e4/torch-2.9.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:2c14b3da5df416cf9cb5efab83aa3056f5b8cd8620b8fde81b4987ecab730587", size = 74480347, upload-time = "2025-11-12T15:21:57.648Z" }, + { url = "https://files.pythonhosted.org/packages/48/50/c4b5112546d0d13cc9eaa1c732b823d676a9f49ae8b6f97772f795874a03/torch-2.9.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1edee27a7c9897f4e0b7c14cfc2f3008c571921134522d5b9b5ec4ebbc69041a", size = 74433245, upload-time = "2025-11-12T15:22:39.027Z" }, + { url = "https://files.pythonhosted.org/packages/81/c9/2628f408f0518b3bae49c95f5af3728b6ab498c8624ab1e03a43dd53d650/torch-2.9.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:19d144d6b3e29921f1fc70503e9f2fc572cde6a5115c0c0de2f7ca8b1483e8b6", size = 104134804, upload-time = "2025-11-12T15:22:35.222Z" }, + { url = "https://files.pythonhosted.org/packages/28/fc/5bc91d6d831ae41bf6e9e6da6468f25330522e92347c9156eb3f1cb95956/torch-2.9.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c432d04376f6d9767a9852ea0def7b47a7bbc8e7af3b16ac9cf9ce02b12851c9", size = 899747132, upload-time = "2025-11-12T15:23:36.068Z" }, + { url = "https://files.pythonhosted.org/packages/63/5d/e8d4e009e52b6b2cf1684bde2a6be157b96fb873732542fb2a9a99e85a83/torch-2.9.1-cp314-cp314-win_amd64.whl", hash = "sha256:d187566a2cdc726fc80138c3cdb260970fab1c27e99f85452721f7759bbd554d", size = 110934845, upload-time = "2025-11-12T15:22:48.367Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b2/2d15a52516b2ea3f414643b8de68fa4cb220d3877ac8b1028c83dc8ca1c4/torch-2.9.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cb10896a1f7fedaddbccc2017ce6ca9ecaaf990f0973bdfcf405439750118d2c", size = 74823558, upload-time = "2025-11-12T15:22:43.392Z" }, + { url = "https://files.pythonhosted.org/packages/86/5c/5b2e5d84f5b9850cd1e71af07524d8cbb74cba19379800f1f9f7c997fc70/torch-2.9.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0a2bd769944991c74acf0c4ef23603b9c777fdf7637f115605a4b2d8023110c7", size = 104145788, upload-time = "2025-11-12T15:23:52.109Z" }, + { url = "https://files.pythonhosted.org/packages/a9/8c/3da60787bcf70add986c4ad485993026ac0ca74f2fc21410bc4eb1bb7695/torch-2.9.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:07c8a9660bc9414c39cac530ac83b1fb1b679d7155824144a40a54f4a47bfa73", size = 899735500, upload-time = "2025-11-12T15:24:08.788Z" }, + { url = "https://files.pythonhosted.org/packages/db/2b/f7818f6ec88758dfd21da46b6cd46af9d1b3433e53ddbb19ad1e0da17f9b/torch-2.9.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c88d3299ddeb2b35dcc31753305612db485ab6f1823e37fb29451c8b2732b87e", size = 111163659, upload-time = "2025-11-12T15:23:20.009Z" }, +] + +[[package]] +name = "torch" +version = "2.9.1+cpu" source = { registry = "https://download.pytorch.org/whl/cpu" } resolution-markers = [ "python_full_version >= '3.12' and sys_platform == 'linux'", @@ -1910,30 +1969,92 @@ dependencies = [ { name = "typing-extensions", marker = "(sys_platform != 'darwin' and extra == 'extra-8-nanochat-cpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, ] wheels = [ - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b224792ea567b52c7f1ce1d789567f6920e06fd3b339fa1e1b05948845f783ad" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:bd2a257e670ede9fc01c6d76dccdc473040913b8e9328169bf177dbdc38e2484" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp310-cp310-win_amd64.whl", hash = "sha256:96f3f7aa4eb9e7fc5af8a722eaf1e5e32e3039dbafe817178d7b90a8566be32d" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:da77341ccaba31762d9238b0942c165c4582a26818f3045b052b39cebdd7ad9d" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:add3e93ecc1eeaa6853f6a973ce60ffb3cb14ed2e80f5055e139b09385dce0a7" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-win_amd64.whl", hash = "sha256:389e1e0b8083fd355f7caf5ba82356b5e01c318998bd575dbf2285a0d8137089" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-win_arm64.whl", hash = "sha256:5ce3d01aef91dc078fbb121814e556d55bc886d303efaf42c4fe67e411f5f9ad" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3a651434ae1248b0568c12b5f9e3acc8942eb28378d9d04a79302938b68c6f24" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:28f6eb31b08180a5c5e98d5bc14eef6909c9f5a1dbff9632c3e02a8773449349" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp312-cp312-win_amd64.whl", hash = "sha256:e438061b87ec7dd6018fca9f975219889aa0a3f6cdc3ea10dd0ae2bc7f1c47ce" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp312-cp312-win_arm64.whl", hash = "sha256:eb13ff1c34e338d722e76a4fd83b8d282782505bd1b99af4b3c32da66eba6eb4" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:be4438d8dad7f0d5a5e54f0feef8a893446894ec87f102bb1d82dcc4518542e4" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6c9b217584400963d5b4daddb3711ec7a3778eab211e18654fba076cce3b8682" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313-win_amd64.whl", hash = "sha256:728372e3f58c5826445f677746e5311c1935c1a7c59599f73a49ded850e038e8" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313-win_arm64.whl", hash = "sha256:95e56c26f919fbb98f16e7a0b87af494b893f9da9a65a020f17a01c13e520a81" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:6c777160288b08555820781ae0f3a2c67a59bd24b065e88ca1ec20e2f9dc8ac7" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:528fd338311f31c9fb18038cafd00e6eae0bf5ad5577521701acb62510753d18" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp313-cp313t-win_amd64.whl", hash = "sha256:d572863990e7d2762b547735ef589f6350d9eb4e441d38753a1c33636698cf4c" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:44aadb735774d4a99525d2ec29126b23016c44a07b02ce6c237dfa61a223dd52" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b355e07b7f0c369cb031adfcbff5c37a609abcea091b918a39886412afd2e07d" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp314-cp314-win_amd64.whl", hash = "sha256:c2698999361d73c2d25d7cc8a787130188d49b183abb18b554228daa102e1594" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:fa0d1373d04b30ff8f12d542135d292f1a1ddb7c0d852a3d487a320360e5dab9" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:2f49bb57a5fe0dc7f8e73ea9e5d36ebda2ea25b8a714a788f0fc2fc47d20a830" }, - { url = "https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp314-cp314t-win_amd64.whl", hash = "sha256:3a60d1ecf27a9cce839b3aa665b26f0af1b1007b9c9f1e7f597f6b7bdf107617" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp310-cp310-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp310-cp310-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp310-cp310-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp311-cp311-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp311-cp311-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp311-cp311-win_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp312-cp312-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp312-cp312-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp312-cp312-win_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313-win_arm64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313t-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313t-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp313-cp313t-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp314-cp314-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp314-cp314-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp314-cp314-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp314-cp314t-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp314-cp314t-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cpu/torch-2.9.1%2Bcpu-cp314-cp314t-win_amd64.whl" }, +] + +[[package]] +name = "torch" +version = "2.9.1+cu128" +source = { registry = "https://download.pytorch.org/whl/cu128" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform == 'linux'", + "python_full_version >= '3.12' and sys_platform != 'linux'", + "python_full_version == '3.11.*' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and sys_platform != 'linux'", + "python_full_version < '3.11' and sys_platform != 'linux'", +] +dependencies = [ + { name = "filelock", marker = "extra == 'extra-8-nanochat-gpu'" }, + { name = "fsspec", marker = "extra == 'extra-8-nanochat-gpu'" }, + { name = "jinja2", marker = "extra == 'extra-8-nanochat-gpu'" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cublas-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cuda-cupti-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cuda-runtime-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cudnn-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cufft-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cufile-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-curand-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusolver-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusparse-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-cusparselt-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nccl-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvshmem-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "nvidia-nvtx-cu12", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "sympy", marker = "extra == 'extra-8-nanochat-gpu'" }, + { name = "triton", version = "3.5.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "typing-extensions", marker = "extra == 'extra-8-nanochat-gpu'" }, +] +wheels = [ + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp310-cp310-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp310-cp310-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp311-cp311-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp311-cp311-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp311-cp311-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp312-cp312-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp312-cp312-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp312-cp312-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp313-cp313-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp313-cp313-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp313-cp313-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp313-cp313t-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp313-cp313t-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp313-cp313t-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp314-cp314-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp314-cp314-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp314-cp314-win_amd64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp314-cp314t-manylinux_2_28_aarch64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp314-cp314t-manylinux_2_28_x86_64.whl" }, + { url = "https://download.pytorch.org/whl/cu128/torch-2.9.1%2Bcu128-cp314-cp314t-win_amd64.whl" }, ] [[package]] @@ -1949,18 +2070,76 @@ wheels = [ ] [[package]] -name = "triton" -version = "3.4.0" +name = "transformers" +version = "4.57.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "setuptools", marker = "(sys_platform == 'linux' and extra == 'extra-8-nanochat-gpu') or (extra == 'extra-8-nanochat-cpu' and extra == 'extra-8-nanochat-gpu')" }, + { name = "filelock" }, + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "requests" }, + { name = "safetensors" }, + { name = "tokenizers" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dd/70/d42a739e8dfde3d92bb2fff5819cbf331fe9657323221e79415cd5eb65ee/transformers-4.57.3.tar.gz", hash = "sha256:df4945029aaddd7c09eec5cad851f30662f8bd1746721b34cc031d70c65afebc", size = 10139680, upload-time = "2025-11-25T15:51:30.139Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/6b/2f416568b3c4c91c96e5a365d164f8a4a4a88030aa8ab4644181fdadce97/transformers-4.57.3-py3-none-any.whl", hash = "sha256:c77d353a4851b1880191603d36acb313411d3577f6e2897814f333841f7003f4", size = 11993463, upload-time = "2025-11-25T15:51:26.493Z" }, +] + +[[package]] +name = "triton" +version = "3.5.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'linux'", ] wheels = [ - { url = "https://files.pythonhosted.org/packages/62/ee/0ee5f64a87eeda19bbad9bc54ae5ca5b98186ed00055281fd40fb4beb10e/triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128", size = 155430069, upload-time = "2025-07-30T19:58:21.715Z" }, - { url = "https://files.pythonhosted.org/packages/7d/39/43325b3b651d50187e591eefa22e236b2981afcebaefd4f2fc0ea99df191/triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b70f5e6a41e52e48cfc087436c8a28c17ff98db369447bcaff3b887a3ab4467", size = 155531138, upload-time = "2025-07-30T19:58:29.908Z" }, - { url = "https://files.pythonhosted.org/packages/d0/66/b1eb52839f563623d185f0927eb3530ee4d5ffe9d377cdaf5346b306689e/triton-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31c1d84a5c0ec2c0f8e8a072d7fd150cab84a9c239eaddc6706c081bfae4eb04", size = 155560068, upload-time = "2025-07-30T19:58:37.081Z" }, - { url = "https://files.pythonhosted.org/packages/30/7b/0a685684ed5322d2af0bddefed7906674f67974aa88b0fae6e82e3b766f6/triton-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00be2964616f4c619193cb0d1b29a99bd4b001d7dc333816073f92cf2a8ccdeb", size = 155569223, upload-time = "2025-07-30T19:58:44.017Z" }, - { url = "https://files.pythonhosted.org/packages/20/63/8cb444ad5cdb25d999b7d647abac25af0ee37d292afc009940c05b82dda0/triton-3.4.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7936b18a3499ed62059414d7df563e6c163c5e16c3773678a3ee3d417865035d", size = 155659780, upload-time = "2025-07-30T19:58:51.171Z" }, + { url = "https://files.pythonhosted.org/packages/dd/22/507b6f58a35e05e84381630b2dc2a3cee1a7a2a7eaf4cba857c638a18a24/triton-3.5.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6f90de6a6566bb619b4c0adc9855729e1b1b5e26533fca1bf6206e96b6d277a3", size = 159827599, upload-time = "2025-10-15T19:15:43.87Z" }, + { url = "https://files.pythonhosted.org/packages/0b/eb/09e31d107a5d00eb281aa7e6635ca463e9bca86515944e399480eadb71f8/triton-3.5.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d5d3b3d480debf24eaa739623c9a42446b0b77f95593d30eb1f64cd2278cc1f0", size = 170333110, upload-time = "2025-10-13T16:37:49.588Z" }, + { url = "https://files.pythonhosted.org/packages/79/f9/b6f60f978397c616fd8dacca2305759fe4f80d397b20ef72534803244bd5/triton-3.5.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8457b22148defefdcb7fa8144b05ce211b9faefad650a1ce85b23df488d5549c", size = 159926731, upload-time = "2025-10-15T19:15:49.682Z" }, + { url = "https://files.pythonhosted.org/packages/3d/78/949a04391c21956c816523678f0e5fa308eb5b1e7622d88c4e4ef5fceca0/triton-3.5.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f34bfa21c5b3a203c0f0eab28dcc1e49bd1f67d22724e77fb6665a659200a4ec", size = 170433488, upload-time = "2025-10-13T16:37:57.132Z" }, + { url = "https://files.pythonhosted.org/packages/87/9b/30988039e1e84df7554fba24e6a734d2d0e847af33cabdf9b532b3c51456/triton-3.5.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7da21fccceafc163e3a5e857abe34351ef76345af06cabf9637a914742671f0b", size = 159946647, upload-time = "2025-10-15T19:15:56.325Z" }, + { url = "https://files.pythonhosted.org/packages/f5/3a/e991574f3102147b642e49637e0281e9bb7c4ba254edb2bab78247c85e01/triton-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c9e71db82261c4ffa3921cd050cd5faa18322d2d405c30eb56084afaff3b0833", size = 170476535, upload-time = "2025-10-13T16:38:05.18Z" }, + { url = "https://files.pythonhosted.org/packages/cd/85/e37f1197acb04c8f3d83851d23d5d6ed5060ef74580668b112e23fdfa203/triton-3.5.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:188da5b81fa2f8322c27fec1627703eac24cb9bb7ab0dfbe9925973bc1b070d3", size = 159958970, upload-time = "2025-10-15T19:16:01.717Z" }, + { url = "https://files.pythonhosted.org/packages/6c/29/10728de8a6e932e517c10773486b8e99f85d1b1d9dd87d9a9616e1fef4a1/triton-3.5.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e6bb9aa5519c084a333acdba443789e50012a4b851cd486c54f0b8dc2a8d3a12", size = 170487289, upload-time = "2025-10-13T16:38:11.662Z" }, + { url = "https://files.pythonhosted.org/packages/b8/1d/38258f05010ac17a7b058c022911c9cae6526e149b7397134a048cf5a6c2/triton-3.5.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:03127d9b33aaf979c856676b394bc059ec1d68cb6da68ae03f62dd8ad77a04ae", size = 160073012, upload-time = "2025-10-15T19:16:07.477Z" }, + { url = "https://files.pythonhosted.org/packages/5c/38/db80e48b9220c9bce872b0f616ad0446cdf554a40b85c7865cbca99ab3c2/triton-3.5.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c83f2343e1a220a716c7b3ab9fccfcbe3ad4020d189549200e2d2e8d5868bed9", size = 170577179, upload-time = "2025-10-13T16:38:17.865Z" }, + { url = "https://files.pythonhosted.org/packages/91/fe/8f5771d00227f4eb1ee034f218ed427102b989366d2275fe3b3c105a3921/triton-3.5.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:468936651d383f4a6d10068d34a627505e13af55be5d002b9f27b987e7a5f0ac", size = 159957460, upload-time = "2025-10-15T19:16:12.626Z" }, + { url = "https://files.pythonhosted.org/packages/ff/60/1810655d1d856c9a4fcc90ee8966d85f552d98c53a6589f95ab2cbe27bb8/triton-3.5.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:da0fa67ccd76c3dcfb0bffe1b1c57c685136a6bd33d141c24d9655d4185b1289", size = 170487949, upload-time = "2025-10-13T16:38:24.881Z" }, + { url = "https://files.pythonhosted.org/packages/78/59/99edd103958fe6e42b50b9ad8ce4f223ddf4ccf475259cf7d2b53381dc6c/triton-3.5.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7ceef21410229ac23173a28eee5cfc0e37c1dfdb8b4bc11ecda2e3ecec7c686", size = 160075629, upload-time = "2025-10-15T19:16:18.746Z" }, + { url = "https://files.pythonhosted.org/packages/fb/b7/1dec8433ac604c061173d0589d99217fe7bf90a70bdc375e745d044b8aad/triton-3.5.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:317fe477ea8fd4524a6a8c499fb0a36984a56d0b75bf9c9cb6133a1c56d5a6e7", size = 170580176, upload-time = "2025-10-13T16:38:31.14Z" }, +] + +[[package]] +name = "triton" +version = "3.5.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.12' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'linux'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/2e/f95e673222afa2c7f0c687d8913e98fcf2589ef0b1405de76894e37fe18f/triton-3.5.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f63e34dcb32d7bd3a1d0195f60f30d2aee8b08a69a0424189b71017e23dfc3d2", size = 159821655, upload-time = "2025-11-11T17:51:44.09Z" }, + { url = "https://files.pythonhosted.org/packages/fd/6e/676ab5019b4dde8b9b7bab71245102fc02778ef3df48218b298686b9ffd6/triton-3.5.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5fc53d849f879911ea13f4a877243afc513187bc7ee92d1f2c0f1ba3169e3c94", size = 170320692, upload-time = "2025-11-11T17:40:46.074Z" }, + { url = "https://files.pythonhosted.org/packages/dc/dc/6ce44d055f2fc2403c4ec6b3cfd3a9b25f57b7d95efadccdea91497f8e81/triton-3.5.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da47169e30a779bade679ce78df4810fca6d78a955843d2ddb11f226adc517dc", size = 159928005, upload-time = "2025-11-11T17:51:50.008Z" }, + { url = "https://files.pythonhosted.org/packages/b0/72/ec90c3519eaf168f22cb1757ad412f3a2add4782ad3a92861c9ad135d886/triton-3.5.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:61413522a48add32302353fdbaaf92daaaab06f6b5e3229940d21b5207f47579", size = 170425802, upload-time = "2025-11-11T17:40:53.209Z" }, + { url = "https://files.pythonhosted.org/packages/db/53/2bcc46879910991f09c063eea07627baef2bc62fe725302ba8f46a2c1ae5/triton-3.5.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:275a045b6ed670dd1bd005c3e6c2d61846c74c66f4512d6f33cc027b11de8fd4", size = 159940689, upload-time = "2025-11-11T17:51:55.938Z" }, + { url = "https://files.pythonhosted.org/packages/f2/50/9a8358d3ef58162c0a415d173cfb45b67de60176e1024f71fbc4d24c0b6d/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d2c6b915a03888ab931a9fd3e55ba36785e1fe70cbea0b40c6ef93b20fc85232", size = 170470207, upload-time = "2025-11-11T17:41:00.253Z" }, + { url = "https://files.pythonhosted.org/packages/f1/ba/805684a992ee32d486b7948d36aed2f5e3c643fc63883bf8bdca1c3f3980/triton-3.5.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:56765ffe12c554cd560698398b8a268db1f616c120007bfd8829d27139abd24a", size = 159955460, upload-time = "2025-11-11T17:52:01.861Z" }, + { url = "https://files.pythonhosted.org/packages/27/46/8c3bbb5b0a19313f50edcaa363b599e5a1a5ac9683ead82b9b80fe497c8d/triton-3.5.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3f4346b6ebbd4fad18773f5ba839114f4826037c9f2f34e0148894cd5dd3dba", size = 170470410, upload-time = "2025-11-11T17:41:06.319Z" }, + { url = "https://files.pythonhosted.org/packages/84/1e/7df59baef41931e21159371c481c31a517ff4c2517343b62503d0cd2be99/triton-3.5.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:02c770856f5e407d24d28ddc66e33cf026e6f4d360dcb8b2fabe6ea1fc758621", size = 160072799, upload-time = "2025-11-11T17:52:07.293Z" }, + { url = "https://files.pythonhosted.org/packages/37/92/e97fcc6b2c27cdb87ce5ee063d77f8f26f19f06916aa680464c8104ef0f6/triton-3.5.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0b4d2c70127fca6a23e247f9348b8adde979d2e7a20391bfbabaac6aebc7e6a8", size = 170579924, upload-time = "2025-11-11T17:41:12.455Z" }, + { url = "https://files.pythonhosted.org/packages/14/f9/0430e879c1e63a1016cb843261528fd3187c872c3a9539132efc39514753/triton-3.5.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f617aa7925f9ea9968ec2e1adaf93e87864ff51549c8f04ce658f29bbdb71e2d", size = 159956163, upload-time = "2025-11-11T17:52:12.999Z" }, + { url = "https://files.pythonhosted.org/packages/a4/e6/c595c35e5c50c4bc56a7bac96493dad321e9e29b953b526bbbe20f9911d0/triton-3.5.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0637b1efb1db599a8e9dc960d53ab6e4637db7d4ab6630a0974705d77b14b60", size = 170480488, upload-time = "2025-11-11T17:41:18.222Z" }, + { url = "https://files.pythonhosted.org/packages/41/1e/63d367c576c75919e268e4fbc33c1cb33b6dc12bb85e8bfe531c2a8bd5d3/triton-3.5.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8932391d7f93698dfe5bc9bead77c47a24f97329e9f20c10786bb230a9083f56", size = 160073620, upload-time = "2025-11-11T17:52:18.403Z" }, + { url = "https://files.pythonhosted.org/packages/16/b5/b0d3d8b901b6a04ca38df5e24c27e53afb15b93624d7fd7d658c7cd9352a/triton-3.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bac7f7d959ad0f48c0e97d6643a1cc0fd5786fe61cb1f83b537c6b2d54776478", size = 170582192, upload-time = "2025-11-11T17:41:23.963Z" }, ] [[package]]