Bobbie/rename torchspec to aurora

.isort.cfg

@@ -1,4 +1,4 @@
 [settings]
 profile=black
-known_first_party=torchspec
+known_first_party=aurora
 filter_files=true

CONTRIBUTING.md

@@ -7,7 +7,7 @@ Thank you for your interest in contributing to Aurora! This guide will help you
 1. Clone the repository and create the conda environment:
 
 ```bash
-git clone https://github.com/torchspec-project/aurora.git
+git clone https://github.com/aurora-project/aurora.git
 cd aurora
 ./tools/build_conda.sh
 micromamba activate aurora

README.md

@@ -1,6 +1,6 @@
 # Aurora
 
-Aurora is a unified training-serving system for online speculative decoding. It closes the loop between speculator training and serving by continuously learning a draft model directly from live inference traces — treating online speculator learning as an asynchronous reinforcement-learning problem. Aurora is built on top of [TorchSpec](https://github.com/xwuShirley/torchspec).
+Aurora is a unified training-serving system for online speculative decoding. It closes the loop between speculator training and serving by continuously learning a draft model directly from live inference traces — treating online speculator learning as an asynchronous reinforcement-learning problem. Aurora is built on top of [TorchSpec](https://github.com/xwuShirley/aurora).
 
 Aurora supports **day-0 deployment**: a speculator can be served immediately and rapidly adapted to live traffic, improving system performance while providing immediate utility feedback. Across experiments, Aurora achieves a **1.5x day-0 speedup** on recently released frontier models (e.g., MiniMax-M2.1 and Qwen3-Coder-Next), and adapts effectively to distribution shifts in user traffic, delivering an additional **1.25x speedup** over a well-trained but static speculator on widely used models (e.g., Qwen3).
 
@@ -50,7 +50,7 @@ See [`examples/README.md`](examples/README.md) for the full example catalog, per
 
 ## Production Notes
 
-- The example `run.sh` scripts are **single-node oriented** — they manage their own local Ray cluster. For multi-node or Kubernetes deployments, start Ray manually and invoke `python3 -m torchspec.train_entry` directly. See [docs/ray.md](docs/ray.md).
+- The example `run.sh` scripts are **single-node oriented** — they manage their own local Ray cluster. For multi-node or Kubernetes deployments, start Ray manually and invoke `python3 -m aurora.train_entry` directly. See [docs/ray.md](docs/ray.md).
 - **External with-draft** mode requires a **shared filesystem** between training and the SGLang server for draft weight sync.
 - `online_serving.hidden_states_dtype` must match the serving model's dtype (e.g., set `float16` when serving an FP8 model).
 - Training and inference GPU sets (`CUDA_VISIBLE_DEVICES` vs `SGLANG_GPUS`) **must not overlap**.
@@ -90,7 +90,7 @@ W&B logging is disabled by default (report_to: none). To enable it, set report_t
 Enable verbose logging:
 
 ```bash
-TORCHSPEC_LOG_LEVEL=DEBUG bash examples/qwen3-4b-external-with-draft/run.sh
+AURORA_LOG_LEVEL=DEBUG bash examples/qwen3-4b-external-with-draft/run.sh
 ```
 
 ## Citation

torchspec/__init__.py → aurora/__init__.py

@@ -20,8 +20,8 @@
 
 """TorchSpec - Torch native spec decode training framework."""
 
-from torchspec.models import Eagle3Model
-from torchspec.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel
+from aurora.models import Eagle3Model
+from aurora.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel
 
 __all__ = [
     "Eagle3Model",

torchspec/config/__init__.py → aurora/config/__init__.py

@@ -18,7 +18,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from torchspec.config.train_config import Config, config_to_flat_args, load_config
+from aurora.config.train_config import Config, config_to_flat_args, load_config
 
 __all__ = [
     "Config",

torchspec/config/inference_config.py → aurora/config/inference_config.py

@@ -29,7 +29,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
-from torchspec.config.mooncake_config import MooncakeConfig
+from aurora.config.mooncake_config import MooncakeConfig
 
 
 @dataclass

torchspec/config/mooncake_config.py → aurora/config/mooncake_config.py

@@ -22,7 +22,7 @@
 from dataclasses import dataclass
 from typing import Tuple
 
-from torchspec.transfer.mooncake.helpers import calculate_eagle3_buffer_size
+from aurora.transfer.mooncake.helpers import calculate_eagle3_buffer_size
 
 
 @dataclass
@@ -116,7 +116,7 @@ def from_flat_args(cls, args) -> "MooncakeConfig":
         - local_hostname auto-resolution via RayActor.get_node_ip()
         - Size string parsing (handled automatically by __post_init__)
         """
-        from torchspec.ray.ray_actor import RayActor
+        from aurora.ray.ray_actor import RayActor
 
         master_server_address = getattr(args, "mooncake_master_server_address", None)
         metadata_port = getattr(args, "mooncake_metadata_port", None)

torchspec/config/train_config.py → aurora/config/train_config.py

@@ -27,9 +27,9 @@
 
 from omegaconf import DictConfig, OmegaConf
 
-from torchspec.config.inference_config import InferenceConfig
-from torchspec.data.utils import is_local_data_path
-from torchspec.utils.logging import logger
+from aurora.config.inference_config import InferenceConfig
+from aurora.data.utils import is_local_data_path
+from aurora.utils.logging import logger
 
 
 @dataclass
@@ -57,7 +57,7 @@ class DebugConfig:
     memory_snapshot_dir: str = "."
     memory_snapshot_num_steps: Optional[int] = None
     memory_snapshot_path: str = ""
-    profile_dir_name: Optional[str] = "/tmp/torchspec_profiles"
+    profile_dir_name: Optional[str] = "/tmp/aurora_profiles"
     profile_step_end: int = 0
     profile_step_start: int = 0
     profile_target: list = field(default_factory=lambda: ["train_overall"])

torchspec/controller/__init__.py → aurora/controller/__init__.py

@@ -18,15 +18,15 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from torchspec.controller.inference_manager import AsyncInferenceManager
-from torchspec.controller.loop import run_training_loop
-from torchspec.controller.setup import (
+from aurora.controller.inference_manager import AsyncInferenceManager
+from aurora.controller.loop import run_training_loop
+from aurora.controller.setup import (
     auto_calculate_training_steps,
     build_mooncake_config,
     setup_async_training,
     setup_async_training_with_engines,
 )
-from torchspec.controller.training_controller import AsyncTrainingController
+from aurora.controller.training_controller import AsyncTrainingController
 
 __all__ = [
     "AsyncTrainingController",

torchspec/controller/inference_manager.py → aurora/controller/inference_manager.py

@@ -46,8 +46,8 @@
 import ray
 from ray.exceptions import RayActorError
 
-from torchspec.utils.logging import logger
-from torchspec.utils.types import InferenceInput, InferenceOutput
+from aurora.utils.logging import logger
+from aurora.utils.types import InferenceInput, InferenceOutput
 
 MOONCAKE_BACKPRESSURE_POLL_INTERVAL = 0.5  # seconds
 MOONCAKE_BACKPRESSURE_LOG_INTERVAL = 5.0  # seconds

torchspec/controller/loop.py → aurora/controller/loop.py

@@ -29,11 +29,11 @@
 import wandb
 from tqdm import tqdm
 
-from torchspec.training.checkpoint import (
+from aurora.training.checkpoint import (
     _read_checkpoint_metadata,
     _write_checkpoint_metadata,
 )
-from torchspec.utils.logging import logger
+from aurora.utils.logging import logger
 
 
 def _is_save_interval_step(step: int, interval: int) -> bool:

torchspec/controller/setup.py → aurora/controller/setup.py

@@ -24,13 +24,13 @@
 
 import ray
 
-from torchspec.utils.env import get_torchspec_env_vars
-from torchspec.utils.logging import logger
+from aurora.utils.env import get_aurora_env_vars
+from aurora.utils.logging import logger
 
 
 def build_mooncake_config(args):
     """Build MooncakeConfig from flat args namespace."""
-    from torchspec.config.mooncake_config import MooncakeConfig
+    from aurora.config.mooncake_config import MooncakeConfig
 
     return MooncakeConfig.from_flat_args(args)
 
@@ -50,8 +50,8 @@ def setup_async_training_with_engines(
         inference_engines: List of Ray actor engine handles for distributed generation.
         controller: Optional pre-created AsyncTrainingController. If None, a new one is created.
     """
-    from torchspec.controller.inference_manager import AsyncInferenceManager
-    from torchspec.controller.training_controller import AsyncTrainingController
+    from aurora.controller.inference_manager import AsyncInferenceManager
+    from aurora.controller.training_controller import AsyncTrainingController
 
     dp_size = (
         getattr(args, "dp_size", None) or args.training_num_nodes * args.training_num_gpus_per_node
@@ -73,7 +73,7 @@ def setup_async_training_with_engines(
 
         driver_node_id = ray.get_runtime_context().get_node_id()
         controller = AsyncTrainingController.options(
-            runtime_env={"env_vars": get_torchspec_env_vars()},
+            runtime_env={"env_vars": get_aurora_env_vars()},
             scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=driver_node_id, soft=False),
         ).remote(args, dp_size)
 

torchspec/controller/training_controller.py → aurora/controller/training_controller.py

@@ -57,10 +57,10 @@
 import ray
 from ray.util.queue import Queue
 
-from torchspec.training.data_fetcher import TrainSample
-from torchspec.utils.logging import logger
-from torchspec.utils.memory import estimate_tensor_bytes
-from torchspec.utils.types import InferenceInput, InferenceOutput
+from aurora.training.data_fetcher import TrainSample
+from aurora.utils.logging import logger
+from aurora.utils.memory import estimate_tensor_bytes
+from aurora.utils.types import InferenceInput, InferenceOutput
 
 _estimate_bytes = estimate_tensor_bytes
 
@@ -196,7 +196,7 @@ def add_dataset(self, dataset: list) -> int:
 
     def load_dataset(self, args) -> int:
         """Load and process dataset on the controller, store for epoch reloads, and prime the prompt buffer."""
-        from torchspec.data.dataset import load_conversation_dataset
+        from aurora.data.dataset import load_conversation_dataset
 
         self._stored_dataset = load_conversation_dataset(args)
         if not self._stored_dataset:
@@ -221,7 +221,7 @@ def load_eval_dataset(self, args) -> int:
         if not eval_data_path:
             return 0
 
-        from torchspec.data.dataset import load_conversation_dataset
+        from aurora.data.dataset import load_conversation_dataset
 
         eval_args = copy.copy(args)
         eval_args.train_data_path = eval_data_path
@@ -261,7 +261,7 @@ def compute_vocab_mapping(self, target_vocab_size: int, draft_vocab_size: int) -
             # the full dataset if the source file contains supervised targets.
             # Fall back to prompt-only warmup tokens when the dataset truly only
             # contains user requests.
-            from torchspec.data.dataset import load_conversation_dataset
+            from aurora.data.dataset import load_conversation_dataset
 
             vocab_args = copy.copy(self.args)
             vocab_args.train_with_decode = False
@@ -284,7 +284,7 @@ def compute_vocab_mapping(self, target_vocab_size: int, draft_vocab_size: int) -
                     for sample in dataset
                 ]
 
-        from torchspec.data.preprocessing import generate_vocab_mapping
+        from aurora.data.preprocessing import generate_vocab_mapping
 
         assert dataset is not None, "No stored dataset for vocab mapping"
         assert "input_ids" in dataset[0], (

torchspec/controller/training_external_server.py → aurora/controller/training_external_server.py

@@ -47,8 +47,8 @@ class PushSampleRequest(BaseModel):
 
         @app.post("/push_sample")
         async def push_sample(req: PushSampleRequest):
-            from torchspec.data.utils import serialize_packed_loss_mask
-            from torchspec.utils.types import InferenceOutput
+            from aurora.data.utils import serialize_packed_loss_mask
+            from aurora.utils.types import InferenceOutput
 
             # Convert shape lists back to tuples
             tensor_shapes = {k: tuple(v) for k, v in req.tensor_shapes.items()}

torchspec/data/__init__.py → aurora/data/__init__.py

@@ -18,13 +18,13 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from torchspec.data.dataset import load_conversation_dataset
-from torchspec.data.preprocessing import (
+from aurora.data.dataset import load_conversation_dataset
+from aurora.data.preprocessing import (
     preprocess_conversations,
     process_token_dict_to_mappings,
 )
-from torchspec.data.template import TEMPLATE_REGISTRY, ChatTemplate
-from torchspec.data.utils import (
+from aurora.data.template import TEMPLATE_REGISTRY, ChatTemplate
+from aurora.data.utils import (
     DataCollatorWithPadding,
     deserialize_packed_loss_mask,
     pack_loss_mask,

torchspec/data/dataset.py → aurora/data/dataset.py

@@ -26,17 +26,17 @@
 import torch
 from tqdm import tqdm
 
-from torchspec.data.preprocessing import _normalize_conversation
-from torchspec.data.template import TEMPLATE_REGISTRY
-from torchspec.data.utils import (
+from aurora.data.preprocessing import _normalize_conversation
+from aurora.data.template import TEMPLATE_REGISTRY
+from aurora.data.utils import (
     estimate_row_count,
     extract_media_urls,
     flatten_multimodal_content,
     load_hf_dataset,
     pack_loss_mask,
     serialize_packed_loss_mask,
 )
-from torchspec.utils.logging import logger
+from aurora.utils.logging import logger
 
 _logging.getLogger("transformers_modules").setLevel(_logging.ERROR)
 
@@ -47,8 +47,8 @@ def _init_tokenize_worker(
     tokenizer_path, trust_remote_code, chat_template_name, last_turn_loss_only=False
 ):
     """Initializer for each worker process — loads tokenizer once."""
-    from torchspec.data.preprocessing import preprocess_conversations
-    from torchspec.utils.processing import load_tokenizer
+    from aurora.data.preprocessing import preprocess_conversations
+    from aurora.utils.processing import load_tokenizer
 
     _logging.getLogger("transformers_modules").setLevel(_logging.ERROR)
     _worker_state["tokenizer"] = load_tokenizer(tokenizer_path, trust_remote_code=trust_remote_code)
@@ -85,8 +85,8 @@ def _tokenize_single(args):
 
 
 def _init_format_worker(tokenizer_path, trust_remote_code, chat_template_name):
-    from torchspec.data.parse import create_parser
-    from torchspec.utils.processing import load_tokenizer
+    from aurora.data.parse import create_parser
+    from aurora.utils.processing import load_tokenizer
 
     _logging.getLogger("transformers_modules").setLevel(_logging.ERROR)
     tokenizer = load_tokenizer(tokenizer_path, trust_remote_code=trust_remote_code)

torchspec/data/parse.py → aurora/data/parse.py

@@ -26,7 +26,7 @@
 import torch
 from transformers import PreTrainedTokenizer
 
-from torchspec.data.template import ChatTemplate
+from aurora.data.template import ChatTemplate
 
 if TYPE_CHECKING:
     from typing import Any

torchspec/data/preprocessing.py → aurora/data/preprocessing.py

@@ -40,15 +40,15 @@
     HAS_QWEN_VL_UTILS = False
     process_vision_info = None
 
-from torchspec.data.parse import create_parser
-from torchspec.data.template import TEMPLATE_REGISTRY, ChatTemplate
-from torchspec.data.utils import (
+from aurora.data.parse import create_parser
+from aurora.data.template import TEMPLATE_REGISTRY, ChatTemplate
+from aurora.data.utils import (
     pack_loss_mask,
     serialize_packed_loss_mask,
     unpack_loss_mask,
 )
-from torchspec.utils.logging import logger
-from torchspec.utils.tensor import padding
+from aurora.utils.logging import logger
+from aurora.utils.tensor import padding
 
 # define a type called conversation
 Conversation = List[Dict[str, str]]

torchspec/data/utils.py → aurora/data/utils.py

@@ -27,7 +27,7 @@
 from datasets import IterableDataset, load_dataset
 from huggingface_hub import hf_hub_download, list_repo_files
 
-from torchspec.models.ops.loss_mask import compute_assistant_loss_mask
+from aurora.models.ops.loss_mask import compute_assistant_loss_mask
 
 _LOCAL_DATA_EXTS = frozenset({".json", ".jsonl", ".parquet", ".arrow", ".csv", ".tsv", ".txt"})
 

torchspec/inference/__init__.py → aurora/inference/__init__.py

@@ -18,7 +18,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from torchspec.inference.factory import (
+from aurora.inference.factory import (
     create_inference_engines,
     prepare_inference_engines,
 )

torchspec/inference/engine/__init__.py → aurora/inference/engine/__init__.py

@@ -18,10 +18,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from torchspec.inference.engine.base import InferenceEngine
-from torchspec.inference.engine.hf_engine import HFEngine
-from torchspec.inference.engine.hf_runner import HFRunner
-from torchspec.inference.engine.sgl_engine import SglEngine
+from aurora.inference.engine.base import InferenceEngine
+from aurora.inference.engine.hf_engine import HFEngine
+from aurora.inference.engine.hf_runner import HFRunner
+from aurora.inference.engine.sgl_engine import SglEngine
 
 __all__ = [
     "InferenceEngine",