{"id":50678455,"url":"https://github.com/anon160/roxxel","last_synced_at":"2026-06-15T17:00:35.889Z","repository":{"id":361632177,"uuid":"1253573673","full_name":"anon160/Roxxel","owner":"anon160","description":"Zero-RAM, JAX-Centric Dataloading, Streaming, and Asynchronous Checkpointing Toolkit","archived":false,"fork":false,"pushed_at":"2026-06-08T15:44:03.000Z","size":333,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-06-08T17:32:22.762Z","etag":null,"topics":["jax","ml","training","training-framework"],"latest_commit_sha":null,"homepage":"https://anon160.github.io/Roxxel/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/anon160.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":"SECURITY.md","support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2026-05-29T15:48:58.000Z","updated_at":"2026-06-08T15:44:27.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/anon160/Roxxel","commit_stats":null,"previous_names":["anon160/roxxel"],"tags_count":28,"template":false,"template_full_name":null,"purl":"pkg:github/anon160/Roxxel","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/anon160%2FRoxxel","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/anon160%2FRoxxel/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/anon160%2FRoxxel/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/anon160%2FRoxxel/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/anon160","download_url":"https://codeload.github.com/anon160/Roxxel/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/anon160%2FRoxxel/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":34372130,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-05-26T15:22:16.424Z","status":"online","status_checked_at":"2026-06-15T02:00:07.085Z","response_time":63,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","ml","training","training-framework"],"created_at":"2026-06-08T17:00:33.579Z","updated_at":"2026-06-15T17:00:35.880Z","avatar_url":"https://github.com/anon160.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Roxxel 🚀 \n\n**Zero-RAM JAX Dataloader, Asynchronous Checkpointer \u0026 Logging Trainer for Flax NNX**\n\nRoxxel is a zero-RAM, ultra-lightweight, and exceptionally fast dataset manager and pre-training orchestrator designed specifically for distributed JAX/Flax NNX clusters. \n\nBy utilizing virtualized POSIX memory-mapped dataset sharding, background asynchronous Orbax checkpointing, and thread-safe distributed logging, Roxxel completely does away with heavy, complex, and over-engineered training frameworks.\n\n---\n\n## 🌟 Key Features\n\n- **OS-Level Memory Mapping (`mmap`)**: Maps multi-terabyte datasets directly into virtual memory via the operating system's kernel page cache. Consumes exactly **0 bytes of Python RAM** for dataset storage.\n- **Unified Causal Streaming**: Automatically chunks, shuffles, and loads batches directly onto JAX device layouts (`jax.make_array_from_process_local_data`) using your Named Sharding mesh. Exposes the exact step count (`len(stream)`).\n- **Instant Offset Seeking**: Resumes streaming from any step index in under 1 millisecond using binary offsets—completely skipping the need to execute slow dummy fast-forward loops.\n- **Dynamic Dtype Auto-Detection**: Detects datatypes (e.g. `int32` token IDs, `float32` arrays, or raw text bytes) on compilation, writes them in a backward-compatible format, and decodes them perfectly on read.\n- **Multi-Dataset Blending / Mixing**: Blends a primary dataset with multiple secondary datasets using weight ratios. Automatically handles dataset exhaustion mid-phase by re-normalizing weights on the fly.\n- **Curriculum Schedule Timelines**: Supports dynamic sequence length extension (e.g., shifting from 1K to 32K context windows) and batch size changes at precise training step boundaries.\n- **Topology-Agnostic Checkpointing**: Offloads PyTree serialization asynchronously to background threads via Orbax Checkpoint Manager. Restores parameters natively using abstract templates, meaning model definition changes do not break older saved weights.\n- **Multi-Host TPU/GPU Pod Safety**: Restrictions ensure only Rank 0 writes stdout prints, system logs, and metrics CSV files, avoiding multi-process locking contention and terminal clutter.\n\n---\n\n## 🤖 What Roxxel Handles Automatically (So You Don't Have To)\n\nRoxxel was engineered to remove the friction of writing accelerator-optimized JAX code. It handles the following tasks automatically under the hood:\n\n1. **JIT Train Step Compilation**: On initialization, `Trainer` automatically compiles your JIT training step (`@nnx.jit`). You don't need to write or trace JAX functions manually.\n2. **Loss Unpacking \u0026 Wrapping**: If your loss function returns a tuple/list (like `(loss, aux_data)`) or a dictionary (like `{\"loss\": loss, \"perplexity\": ppl}`), Roxxel automatically extracts the scalar loss for gradients (`nnx.value_and_grad`), avoiding compiler crashes.\n3. **Auto-Initialized ModelState**: You don't have to write state wrapper boilerplate (e.g. `TrainState` classes). The trainer dynamically bundles your model, optimizer, and training step counter.\n4. **Auto-Initialized Checkpointer \u0026 Logger**: If you pass a folder path to the `save_path` parameter, Roxxel automatically initializes the asynchronous Orbax Checkpointer and multi-threaded Logger under that directory, saving checkpoints in `save_path/checkpoints` and logs directly in `save_path`.\n5. **Dynamic Stream Re-instantiation**: During curriculum phase transitions (e.g. when sequence length changes), Roxxel automatically closes the active stream, computes completed offsets, and swaps the dataset streams instantly.\n6. **Exhaustion Re-normalization**: If one of your mixed datasets runs out of records mid-training, Roxxel removes it from the choice pool and re-normalizes the weights of the remaining active datasets to prevent training stalls.\n7. **Crash-Safe Log Flushing**: In the event of an OOM, crash, or exception, the logger intercepts the exception, logs the stack trace to the system log, and flushes all pending file and stdout writes before bubbling the error.\n8. **Asynchronous NaN Loss Guard**: The trainer periodically checks the materialized loss value at logging or checkpointing intervals. If a NaN is detected, it immediately throws a `ValueError` to halt the run and prevent compute budget waste on divergent parameters.\n\n---\n\n## 📦 Installation\n\nTo install Roxxel:\n```bash\npip install roxxel\n```\n\n---\n\n## 🚀 Quick Start Cookbook\n\nHere is a complete, zero-boilerplate example showing how to initialize a model, define a curriculum schedule, and execute pre-training using the consolidated `Trainer`:\n\n```python\nimport jax\nimport optax\nfrom flax import nnx\nfrom roxxel import Roxxel, Curriculum, Trainer\n\n# 1. Initialize Flax NNX model and optimizer\nmodel = nnx.Linear(10, 5, rngs=nnx.Rngs(42))\ntx = optax.sgd(0.01)\noptimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)\n\n# 2. Define the curriculum (e.g., Phase 1: 1000 steps, Phase 2: 500 steps)\nphases = [\n    {\"steps\": 1000, \"batch_size\": 16, \"seq_len\": 128},\n    {\"steps\": 500, \"batch_size\": 4, \"seq_len\": 512}\n]\ncurriculum = Curriculum(\n    primary_streamer=Roxxel(\"./wiki_tokens_*.rox\"), \n    phases=phases\n)\n\n# 3. Define the loss function\ndef loss_fn(model, batch):\n    logits = model(batch[:, :-1].astype(jax.numpy.float32))\n    targets = batch[:, 1:].astype(jax.numpy.float32)\n    return jax.numpy.mean((logits - targets) ** 2)\n\n# 4. Initialize the Trainer\n# Setting save_path automatically initializes Checkpointer, Logger, and ModelState\ntrainer = Trainer(\n    model=model,\n    optimizer=optimizer,\n    curriculum=curriculum,\n    loss_fn=loss_fn,\n    save_path=\"./run_delta\",\n    checkpoint_every=100,\n    log_every=10\n)\n\n# 5. Run curriculum training\ntrainer.run()\n```\n\n---\n\n## 📖 Learn More\n\nFor complete documentation, design guides, and API specs, visit the [Roxxel Documentation Site](https://anon160.github.io/Roxxel/):\n- **[Why Roxxel?](https://anon160.github.io/Roxxel/why_roxxel/)**: Design philosophy and comparisons to PyTorch/TensorFlow dataloaders.\n- **[End-to-End Tutorial](https://anon160.github.io/Roxxel/tutorial/)**: Full cookbook with hardware sharding and dataset blending.\n- **[Dataloader \u0026 Streaming](https://anon160.github.io/Roxxel/dataloader_and_streaming/)**: Deep dive into block virtualization and streams.\n- **[Curriculum Blending](https://anon160.github.io/Roxxel/curriculum/)**: How to mix datasets and define phases.\n- **[Trainer Orchestration](https://anon160.github.io/Roxxel/trainer/)**: Full Trainer and ModelState specifications.\n- **[Asynchronous Checkpointing](https://anon160.github.io/Roxxel/checkpointing/)**: Orbax asynchronous serialization details.\n- **[Asynchronous Logging](https://anon160.github.io/Roxxel/logging/)**: Rank-zero queue writing details.\n\n---\n\n## ⚖️ License\nMIT License. Feel free to use, modify, and distribute.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fanon160%2Froxxel","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fanon160%2Froxxel","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fanon160%2Froxxel/lists"}