diff --git a/.isort.cfg b/.isort.cfg index 88a5980..411c27a 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -1,4 +1,4 @@ [settings] profile=black -known_first_party=torchspec +known_first_party=aurora filter_files=true diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index e1e8c25..2fd0e25 100644 --- a/CONTRIBUTING.md +++ b/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 diff --git a/README.md b/README.md index a384cd3..ad129f5 100644 --- a/README.md +++ b/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 diff --git a/torchspec/__init__.py b/aurora/__init__.py similarity index 91% rename from torchspec/__init__.py rename to aurora/__init__.py index df5122c..01a8e7c 100644 --- a/torchspec/__init__.py +++ b/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", diff --git a/torchspec/config/__init__.py b/aurora/config/__init__.py similarity index 93% rename from torchspec/config/__init__.py rename to aurora/config/__init__.py index fedceaf..d6b9bef 100644 --- a/torchspec/config/__init__.py +++ b/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", diff --git a/torchspec/config/inference_config.py b/aurora/config/inference_config.py similarity index 98% rename from torchspec/config/inference_config.py rename to aurora/config/inference_config.py index e58dcb5..6720e2c 100644 --- a/torchspec/config/inference_config.py +++ b/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 diff --git a/torchspec/config/mooncake_config.py b/aurora/config/mooncake_config.py similarity index 98% rename from torchspec/config/mooncake_config.py rename to aurora/config/mooncake_config.py index 50f45e2..0875a66 100644 --- a/torchspec/config/mooncake_config.py +++ b/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) diff --git a/torchspec/config/train_config.py b/aurora/config/train_config.py similarity index 98% rename from torchspec/config/train_config.py rename to aurora/config/train_config.py index d2d8a91..a4c304c 100644 --- a/torchspec/config/train_config.py +++ b/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"]) diff --git a/torchspec/config/utils.py b/aurora/config/utils.py similarity index 100% rename from torchspec/config/utils.py rename to aurora/config/utils.py diff --git a/torchspec/controller/__init__.py b/aurora/controller/__init__.py similarity index 85% rename from torchspec/controller/__init__.py rename to aurora/controller/__init__.py index ae29f52..190ab0b 100644 --- a/torchspec/controller/__init__.py +++ b/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", diff --git a/torchspec/controller/inference_manager.py b/aurora/controller/inference_manager.py similarity index 99% rename from torchspec/controller/inference_manager.py rename to aurora/controller/inference_manager.py index 6cde310..d13e2c8 100644 --- a/torchspec/controller/inference_manager.py +++ b/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 diff --git a/torchspec/controller/loop.py b/aurora/controller/loop.py similarity index 99% rename from torchspec/controller/loop.py rename to aurora/controller/loop.py index 5f687ee..7e6bbfa 100644 --- a/torchspec/controller/loop.py +++ b/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: diff --git a/torchspec/controller/setup.py b/aurora/controller/setup.py similarity index 94% rename from torchspec/controller/setup.py rename to aurora/controller/setup.py index 34afcbf..3f79f19 100644 --- a/torchspec/controller/setup.py +++ b/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) diff --git a/torchspec/controller/training_controller.py b/aurora/controller/training_controller.py similarity index 98% rename from torchspec/controller/training_controller.py rename to aurora/controller/training_controller.py index 5b94706..07f2825 100644 --- a/torchspec/controller/training_controller.py +++ b/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], ( diff --git a/torchspec/controller/training_external_server.py b/aurora/controller/training_external_server.py similarity index 96% rename from torchspec/controller/training_external_server.py rename to aurora/controller/training_external_server.py index 1c90c52..b9c82ba 100644 --- a/torchspec/controller/training_external_server.py +++ b/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()} diff --git a/torchspec/data/__init__.py b/aurora/data/__init__.py similarity index 88% rename from torchspec/data/__init__.py rename to aurora/data/__init__.py index 531b180..261070b 100644 --- a/torchspec/data/__init__.py +++ b/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, diff --git a/torchspec/data/dataset.py b/aurora/data/dataset.py similarity index 96% rename from torchspec/data/dataset.py rename to aurora/data/dataset.py index 07ffecc..71aeb65 100644 --- a/torchspec/data/dataset.py +++ b/aurora/data/dataset.py @@ -26,9 +26,9 @@ 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, @@ -36,7 +36,7 @@ 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) diff --git a/torchspec/data/parse.py b/aurora/data/parse.py similarity index 99% rename from torchspec/data/parse.py rename to aurora/data/parse.py index 55a952e..123486a 100644 --- a/torchspec/data/parse.py +++ b/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 diff --git a/torchspec/data/preprocessing.py b/aurora/data/preprocessing.py similarity index 98% rename from torchspec/data/preprocessing.py rename to aurora/data/preprocessing.py index 9527ead..0f5bae1 100644 --- a/torchspec/data/preprocessing.py +++ b/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]] diff --git a/torchspec/data/template.py b/aurora/data/template.py similarity index 100% rename from torchspec/data/template.py rename to aurora/data/template.py diff --git a/torchspec/data/utils.py b/aurora/data/utils.py similarity index 99% rename from torchspec/data/utils.py rename to aurora/data/utils.py index bf19d6c..e5b19a2 100644 --- a/torchspec/data/utils.py +++ b/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"}) diff --git a/torchspec/inference/__init__.py b/aurora/inference/__init__.py similarity index 96% rename from torchspec/inference/__init__.py rename to aurora/inference/__init__.py index 84ff98f..30174ce 100644 --- a/torchspec/inference/__init__.py +++ b/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, ) diff --git a/torchspec/inference/engine/__init__.py b/aurora/inference/engine/__init__.py similarity index 83% rename from torchspec/inference/engine/__init__.py rename to aurora/inference/engine/__init__.py index cd0a299..d4dc182 100644 --- a/torchspec/inference/engine/__init__.py +++ b/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", diff --git a/torchspec/inference/engine/base.py b/aurora/inference/engine/base.py similarity index 100% rename from torchspec/inference/engine/base.py rename to aurora/inference/engine/base.py diff --git a/torchspec/inference/engine/hf_engine.py b/aurora/inference/engine/hf_engine.py similarity index 95% rename from torchspec/inference/engine/hf_engine.py rename to aurora/inference/engine/hf_engine.py index c8b5bfd..e30d684 100644 --- a/torchspec/inference/engine/hf_engine.py +++ b/aurora/inference/engine/hf_engine.py @@ -31,9 +31,9 @@ import ray import torch -from torchspec.inference.engine.base import InferenceEngine -from torchspec.ray.ray_actor import RayActor -from torchspec.utils.logging import logger, setup_file_logging +from aurora.inference.engine.base import InferenceEngine +from aurora.ray.ray_actor import RayActor +from aurora.utils.logging import logger, setup_file_logging class HFEngine(InferenceEngine, RayActor): @@ -67,7 +67,7 @@ def init(self, mooncake_config=None) -> None: Args: mooncake_config: MooncakeConfig object for distributed storage. """ - from torchspec.inference.engine.hf_runner import HFRunner + from aurora.inference.engine.hf_runner import HFRunner if self.base_gpu_id is not None: local_gpu_id = self.setup_gpu(self.base_gpu_id) @@ -79,7 +79,7 @@ def init(self, mooncake_config=None) -> None: self._mooncake_config = mooncake_config if mooncake_config is not None: - from torchspec.transfer.mooncake.utils import ( + from aurora.transfer.mooncake.utils import ( check_mooncake_master_available, ) diff --git a/torchspec/inference/engine/hf_runner.py b/aurora/inference/engine/hf_runner.py similarity index 97% rename from torchspec/inference/engine/hf_runner.py rename to aurora/inference/engine/hf_runner.py index 615b701..a2d0e87 100644 --- a/torchspec/inference/engine/hf_runner.py +++ b/aurora/inference/engine/hf_runner.py @@ -33,11 +33,11 @@ import torch import torch.distributed as dist -from torchspec.config.inference_config import HFInferenceConfig -from torchspec.config.mooncake_config import MooncakeConfig -from torchspec.models.target import HFTargetModel -from torchspec.transfer.mooncake.eagle_store import EagleMooncakeStore -from torchspec.utils.logging import logger +from aurora.config.inference_config import HFInferenceConfig +from aurora.config.mooncake_config import MooncakeConfig +from aurora.models.target import HFTargetModel +from aurora.transfer.mooncake.eagle_store import EagleMooncakeStore +from aurora.utils.logging import logger class HFRunner: diff --git a/torchspec/inference/engine/sgl_engine.py b/aurora/inference/engine/sgl_engine.py similarity index 98% rename from torchspec/inference/engine/sgl_engine.py rename to aurora/inference/engine/sgl_engine.py index 331e53e..2ecd47b 100644 --- a/torchspec/inference/engine/sgl_engine.py +++ b/aurora/inference/engine/sgl_engine.py @@ -36,11 +36,11 @@ import torch from omegaconf import DictConfig, OmegaConf -from torchspec.data.utils import serialize_packed_loss_mask -from torchspec.inference.engine.base import InferenceEngine -from torchspec.ray.ray_actor import RayActor -from torchspec.utils.logging import logger, setup_file_logging -from torchspec.utils.misc import get_default_eagle3_aux_layer_ids, get_free_port +from aurora.data.utils import serialize_packed_loss_mask +from aurora.inference.engine.base import InferenceEngine +from aurora.ray.ray_actor import RayActor +from aurora.utils.logging import logger, setup_file_logging +from aurora.utils.misc import get_default_eagle3_aux_layer_ids, get_free_port # Keys that users might plausibly put in extra_args but are managed by # TorchSpec. Used only to emit a warning — the actual protection comes @@ -146,7 +146,7 @@ def init(self, mooncake_config=None, dist_init_addr: str | None = None) -> None: f"device_name={mooncake_config.device_name}" ) - from torchspec.transfer.mooncake.utils import ( + from aurora.transfer.mooncake.utils import ( check_mooncake_master_available, ) diff --git a/torchspec/inference/factory.py b/aurora/inference/factory.py similarity index 97% rename from torchspec/inference/factory.py rename to aurora/inference/factory.py index a0a725d..a36cf4f 100644 --- a/torchspec/inference/factory.py +++ b/aurora/inference/factory.py @@ -23,10 +23,10 @@ import ray from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from torchspec.inference.engine.hf_engine import HFEngine -from torchspec.inference.engine.sgl_engine import SglEngine -from torchspec.utils.env import get_torchspec_env_vars -from torchspec.utils.logging import logger +from aurora.inference.engine.hf_engine import HFEngine +from aurora.inference.engine.sgl_engine import SglEngine +from aurora.utils.env import get_aurora_env_vars +from aurora.utils.logging import logger # Multi-node TP worker engines must stay alive to participate in NCCL # operations but are never called directly. Store refs here to prevent GC. @@ -181,7 +181,7 @@ def _prepare_sgl_engines( pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg SglRayActor = ray.remote(SglEngine) - env_vars = get_torchspec_env_vars() + env_vars = get_aurora_env_vars() # Step 1: Create all engine actors (without calling init yet) engines = [] @@ -310,7 +310,7 @@ def _create_and_init_actors( placement_group_bundle_index=reordered_bundle_indices[i * num_gpus_per_engine], ) - env_vars = get_torchspec_env_vars() + env_vars = get_aurora_env_vars() constructor_kwargs = { "args": args, diff --git a/torchspec/models/__init__.py b/aurora/models/__init__.py similarity index 88% rename from torchspec/models/__init__.py rename to aurora/models/__init__.py index 9742018..d07b1b2 100644 --- a/torchspec/models/__init__.py +++ b/aurora/models/__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.models.eagle3 import ( +from aurora.models.eagle3 import ( Eagle3Model, compute_lazy_target_padded, compute_target_p_padded, ) -from torchspec.models.ops.loss import compiled_forward_kl_loss -from torchspec.models.ops.loss_mask import compute_assistant_loss_mask +from aurora.models.ops.loss import compiled_forward_kl_loss +from aurora.models.ops.loss_mask import compute_assistant_loss_mask __all__ = [ "Eagle3Model", diff --git a/torchspec/models/draft/__init__.py b/aurora/models/draft/__init__.py similarity index 85% rename from torchspec/models/draft/__init__.py rename to aurora/models/draft/__init__.py index d31be32..31222aa 100644 --- a/torchspec/models/draft/__init__.py +++ b/aurora/models/draft/__init__.py @@ -18,9 +18,9 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -from torchspec.models.draft.auto import AutoDraftModelConfig, AutoEagle3DraftModel -from torchspec.models.draft.base import Eagle3DraftModel -from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3 +from aurora.models.draft.auto import AutoDraftModelConfig, AutoEagle3DraftModel +from aurora.models.draft.base import Eagle3DraftModel +from aurora.models.draft.llama3_eagle import LlamaForCausalLMEagle3 __all__ = [ "AutoDraftModelConfig", diff --git a/torchspec/models/draft/auto.py b/aurora/models/draft/auto.py similarity index 96% rename from torchspec/models/draft/auto.py rename to aurora/models/draft/auto.py index 955f51d..75d8b29 100644 --- a/torchspec/models/draft/auto.py +++ b/aurora/models/draft/auto.py @@ -25,8 +25,8 @@ from transformers import AutoModelForCausalLM as AutoModelForCausalLMBase from transformers import LlamaConfig, PretrainedConfig, modeling_utils -from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3 -from torchspec.utils.logging import logger +from aurora.models.draft.llama3_eagle import LlamaForCausalLMEagle3 +from aurora.utils.logging import logger class AutoEagle3DraftModel(AutoModelForCausalLMBase): diff --git a/torchspec/models/draft/base.py b/aurora/models/draft/base.py similarity index 100% rename from torchspec/models/draft/base.py rename to aurora/models/draft/base.py diff --git a/torchspec/models/draft/llama3_eagle.py b/aurora/models/draft/llama3_eagle.py similarity index 99% rename from torchspec/models/draft/llama3_eagle.py rename to aurora/models/draft/llama3_eagle.py index 16a0cef..8b686c5 100644 --- a/torchspec/models/draft/llama3_eagle.py +++ b/aurora/models/draft/llama3_eagle.py @@ -28,13 +28,13 @@ from transformers.activations import ACT2FN from transformers.models.llama.configuration_llama import LlamaConfig -from torchspec.models.draft.base import Eagle3DraftModel -from torchspec.models.ops.flex_attention import ( +from aurora.models.draft.base import Eagle3DraftModel +from aurora.models.ops.flex_attention import ( compile_friendly_create_block_mask, compile_friendly_flex_attention, generate_eagle3_mask, ) -from torchspec.utils.logging import logger, print_with_rank +from aurora.utils.logging import logger, print_with_rank _flash_attn_import_error: ImportError | None = None try: @@ -69,17 +69,17 @@ def _patch_cutlass_compilation() -> None: --opt-level, so the parser default (3) is used. Lower levels reduce LLVM-IR and ptxas work at the cost of slightly slower kernels. - Controlled by TORCHSPEC_FLASH_ATTN_OPT_LEVEL (default: 3): + Controlled by AURORA_FLASH_ATTN_OPT_LEVEL (default: 3): 3 – current default; ~12 min; fastest kernel 2 – ~6–8 min; <5% slower kernel (recommended for training) 1 – ~3–5 min; ~15% slower kernel 0 – ~1–2 min; ~50% slower kernel (debugging only) - TORCHSPEC_FLASH_ATTN_PTXAS_OPT controls ptxas separately + AURORA_FLASH_ATTN_PTXAS_OPT controls ptxas separately (default: 3). Use 1 for additional compile-time savings when combined with opt-level 1: - export TORCHSPEC_FLASH_ATTN_OPT_LEVEL=1 - export TORCHSPEC_FLASH_ATTN_PTXAS_OPT=1 + export AURORA_FLASH_ATTN_OPT_LEVEL=1 + export AURORA_FLASH_ATTN_PTXAS_OPT=1 """ import os @@ -135,8 +135,8 @@ def _compile_and_cache_with_disk( if getattr(CompileCallable, "_opt_level_patched", False): return - _opt_level = int(os.environ.get("TORCHSPEC_FLASH_ATTN_OPT_LEVEL", "3")) - _ptxas_opt = int(os.environ.get("TORCHSPEC_FLASH_ATTN_PTXAS_OPT", "3")) + _opt_level = int(os.environ.get("AURORA_FLASH_ATTN_OPT_LEVEL", "3")) + _ptxas_opt = int(os.environ.get("AURORA_FLASH_ATTN_PTXAS_OPT", "3")) _orig_compile_callable = CompileCallable._compile @@ -156,7 +156,7 @@ def _compile_with_opt_level(self, func, *args, **kwargs): logger.debug( f"flash_attn compilation: opt-level={_opt_level}, " f"ptxas-opt={_ptxas_opt} " - f"(TORCHSPEC_FLASH_ATTN_OPT_LEVEL / TORCHSPEC_FLASH_ATTN_PTXAS_OPT)" + f"(AURORA_FLASH_ATTN_OPT_LEVEL / AURORA_FLASH_ATTN_PTXAS_OPT)" ) _patch_cutlass_compilation() diff --git a/torchspec/models/eagle3.py b/aurora/models/eagle3.py similarity index 99% rename from torchspec/models/eagle3.py rename to aurora/models/eagle3.py index acacfc2..a538c1d 100644 --- a/torchspec/models/eagle3.py +++ b/aurora/models/eagle3.py @@ -26,11 +26,11 @@ import torch.nn.functional as F from torch.utils.checkpoint import checkpoint as torch_checkpoint -from torchspec.models.ops.loss import ( +from aurora.models.ops.loss import ( compiled_forward_kl_loss, compiled_forward_kl_loss_from_hs, ) -from torchspec.utils.tensor import padding +from aurora.utils.tensor import padding @dataclass diff --git a/torchspec/models/ops/__init__.py b/aurora/models/ops/__init__.py similarity index 88% rename from torchspec/models/ops/__init__.py rename to aurora/models/ops/__init__.py index 10b10d4..a51e833 100644 --- a/torchspec/models/ops/__init__.py +++ b/aurora/models/ops/__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.models.ops.flex_attention import ( +from aurora.models.ops.flex_attention import ( compile_friendly_create_block_mask, compile_friendly_flex_attention, generate_eagle3_mask, ) -from torchspec.models.ops.loss import compiled_forward_kl_loss -from torchspec.models.ops.loss_mask import compute_assistant_loss_mask +from aurora.models.ops.loss import compiled_forward_kl_loss +from aurora.models.ops.loss_mask import compute_assistant_loss_mask __all__ = [ "compile_friendly_create_block_mask", diff --git a/torchspec/models/ops/flex_attention.py b/aurora/models/ops/flex_attention.py similarity index 100% rename from torchspec/models/ops/flex_attention.py rename to aurora/models/ops/flex_attention.py diff --git a/torchspec/models/ops/loss.py b/aurora/models/ops/loss.py similarity index 100% rename from torchspec/models/ops/loss.py rename to aurora/models/ops/loss.py diff --git a/torchspec/models/ops/loss_mask.py b/aurora/models/ops/loss_mask.py similarity index 100% rename from torchspec/models/ops/loss_mask.py rename to aurora/models/ops/loss_mask.py diff --git a/torchspec/models/target/__init__.py b/aurora/models/target/__init__.py similarity index 91% rename from torchspec/models/target/__init__.py rename to aurora/models/target/__init__.py index 7a85cd3..c3f4219 100644 --- a/torchspec/models/target/__init__.py +++ b/aurora/models/target/__init__.py @@ -18,12 +18,12 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -from torchspec.models.target.eagle3_target_model import ( +from aurora.models.target.eagle3_target_model import ( Eagle3TargetModel, Eagle3TargetOutput, HFTargetModel, ) -from torchspec.models.target.target_utils import TargetLMHead +from aurora.models.target.target_utils import TargetLMHead __all__ = [ "Eagle3TargetModel", diff --git a/torchspec/models/target/eagle3_target_model.py b/aurora/models/target/eagle3_target_model.py similarity index 99% rename from torchspec/models/target/eagle3_target_model.py rename to aurora/models/target/eagle3_target_model.py index f11f86d..71a29af 100644 --- a/torchspec/models/target/eagle3_target_model.py +++ b/aurora/models/target/eagle3_target_model.py @@ -26,7 +26,7 @@ import torch.nn as nn from transformers import AutoModelForCausalLM -from torchspec.utils.distributed import get_tp_device_mesh, get_tp_group +from aurora.utils.distributed import get_tp_device_mesh, get_tp_group @dataclass diff --git a/torchspec/models/target/target_utils.py b/aurora/models/target/target_utils.py similarity index 100% rename from torchspec/models/target/target_utils.py rename to aurora/models/target/target_utils.py diff --git a/torchspec/ray/__init__.py b/aurora/ray/__init__.py similarity index 100% rename from torchspec/ray/__init__.py rename to aurora/ray/__init__.py diff --git a/torchspec/ray/placement_group.py b/aurora/ray/placement_group.py similarity index 99% rename from torchspec/ray/placement_group.py rename to aurora/ray/placement_group.py index 2d10dfb..15800f4 100644 --- a/torchspec/ray/placement_group.py +++ b/aurora/ray/placement_group.py @@ -26,8 +26,8 @@ from ray.util.placement_group import placement_group from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from torchspec.ray.train_group import RayTrainGroup -from torchspec.utils.logging import logger +from aurora.ray.train_group import RayTrainGroup +from aurora.utils.logging import logger @ray.remote(num_gpus=1) diff --git a/torchspec/ray/ray_actor.py b/aurora/ray/ray_actor.py similarity index 95% rename from torchspec/ray/ray_actor.py rename to aurora/ray/ray_actor.py index cfb88c5..3d62be6 100644 --- a/torchspec/ray/ray_actor.py +++ b/aurora/ray/ray_actor.py @@ -25,8 +25,8 @@ import torch from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy -from torchspec.utils.logging import logger -from torchspec.utils.misc import _to_local_gpu_id, get_current_node_ip, get_free_port +from aurora.utils.logging import logger +from aurora.utils.misc import _to_local_gpu_id, get_current_node_ip, get_free_port def node_affinity_for_ip(ip: str, name: str = None) -> NodeAffinitySchedulingStrategy: @@ -54,7 +54,7 @@ def node_affinity_for_ip(ip: str, name: str = None) -> NodeAffinitySchedulingStr class RayActor: - """Base class for all torchspec Ray actors.""" + """Base class for all aurora Ray actors.""" @staticmethod def get_node_ip() -> str: diff --git a/torchspec/ray/train_group.py b/aurora/ray/train_group.py similarity index 98% rename from torchspec/ray/train_group.py rename to aurora/ray/train_group.py index 8743199..8bdf631 100644 --- a/torchspec/ray/train_group.py +++ b/aurora/ray/train_group.py @@ -26,7 +26,7 @@ from ray.util.placement_group import PlacementGroup from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from torchspec.utils.env import get_torchspec_env_vars +from aurora.utils.env import get_aurora_env_vars class RayTrainGroup: @@ -80,7 +80,7 @@ def _allocate_gpus_for_training(self, pg, num_gpus_per_actor): train_env_vars = json.loads(train_env_vars) if train_env_vars else {} env_vars = { - **get_torchspec_env_vars(), + **get_aurora_env_vars(), "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": os.environ.get( "NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1" diff --git a/torchspec/train_entry.py b/aurora/train_entry.py similarity index 95% rename from torchspec/train_entry.py rename to aurora/train_entry.py index 03e650f..7a1ef28 100644 --- a/torchspec/train_entry.py +++ b/aurora/train_entry.py @@ -31,25 +31,25 @@ from omegaconf import OmegaConf from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy -from torchspec import AutoDraftModelConfig -from torchspec.config.train_config import config_to_flat_args, load_config -from torchspec.config.utils import generate_draft_model_config -from torchspec.controller import ( +from aurora import AutoDraftModelConfig +from aurora.config.train_config import config_to_flat_args, load_config +from aurora.config.utils import generate_draft_model_config +from aurora.controller import ( AsyncTrainingController, auto_calculate_training_steps, build_mooncake_config, run_training_loop, setup_async_training_with_engines, ) -from torchspec.inference import prepare_inference_engines -from torchspec.ray.placement_group import ( +from aurora.inference import prepare_inference_engines +from aurora.ray.placement_group import ( allocate_train_group, create_placement_groups, ) -from torchspec.training.trainer_actor import TrainerActor -from torchspec.transfer.mooncake.utils import launch_mooncake_master -from torchspec.utils.env import get_torchspec_env_vars -from torchspec.utils.logging import init_tracking, logger +from aurora.training.trainer_actor import TrainerActor +from aurora.transfer.mooncake.utils import launch_mooncake_master +from aurora.utils.env import get_aurora_env_vars +from aurora.utils.logging import init_tracking, logger _Phase = namedtuple("_Phase", ["name", "duration", "is_async", "blocked"]) @@ -204,7 +204,7 @@ def train_async_no_generation(args): with timer.phase("Create controller"): 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, args.dp_size) @@ -344,7 +344,7 @@ def train_async_no_generation(args): # [9.5] Start external training server if online serving is enabled training_server = None if getattr(args, "online_serving_enabled", False): - from torchspec.controller.training_external_server import TrainingExternalServer + from aurora.controller.training_external_server import TrainingExternalServer with timer.phase("Start training external server"): training_server = TrainingExternalServer.remote(args, controller) diff --git a/torchspec/training/__init__.py b/aurora/training/__init__.py similarity index 100% rename from torchspec/training/__init__.py rename to aurora/training/__init__.py diff --git a/torchspec/training/checkpoint.py b/aurora/training/checkpoint.py similarity index 99% rename from torchspec/training/checkpoint.py rename to aurora/training/checkpoint.py index b4a5175..c70b032 100644 --- a/torchspec/training/checkpoint.py +++ b/aurora/training/checkpoint.py @@ -31,7 +31,7 @@ from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict from torch.distributed.checkpoint.stateful import Stateful -from torchspec.utils.logging import logger +from aurora.utils.logging import logger class ModelState(Stateful): diff --git a/torchspec/training/data_fetcher.py b/aurora/training/data_fetcher.py similarity index 99% rename from torchspec/training/data_fetcher.py rename to aurora/training/data_fetcher.py index 1c42d7a..581bd03 100644 --- a/torchspec/training/data_fetcher.py +++ b/aurora/training/data_fetcher.py @@ -32,7 +32,7 @@ from ray.util.queue import Queue as RayQueue from torch.utils.data import DataLoader, IterableDataset -from torchspec.utils.logging import logger +from aurora.utils.logging import logger @dataclass diff --git a/torchspec/training/eagle3_trainer.py b/aurora/training/eagle3_trainer.py similarity index 95% rename from torchspec/training/eagle3_trainer.py rename to aurora/training/eagle3_trainer.py index 3ee3de7..6ba4247 100644 --- a/torchspec/training/eagle3_trainer.py +++ b/aurora/training/eagle3_trainer.py @@ -24,16 +24,16 @@ import torch import torch.distributed as dist -from torchspec import AutoDraftModelConfig, AutoEagle3DraftModel, Eagle3Model -from torchspec.models.eagle3 import compute_lazy_target_padded, compute_target_p_padded -from torchspec.training import checkpoint -from torchspec.training.fsdp import apply_fsdp2, fsdp2_load_full_state_dict -from torchspec.training.optimizer import BF16Optimizer -from torchspec.training.trainer import Trainer -from torchspec.utils.distributed import get_gloo_group -from torchspec.utils.logging import logger -from torchspec.utils.tensor import padding -from torchspec.utils.train_dump import dump_eagle3_batch +from aurora import AutoDraftModelConfig, AutoEagle3DraftModel, Eagle3Model +from aurora.models.eagle3 import compute_lazy_target_padded, compute_target_p_padded +from aurora.training import checkpoint +from aurora.training.fsdp import apply_fsdp2, fsdp2_load_full_state_dict +from aurora.training.optimizer import BF16Optimizer +from aurora.training.trainer import Trainer +from aurora.utils.distributed import get_gloo_group +from aurora.utils.logging import logger +from aurora.utils.tensor import padding +from aurora.utils.train_dump import dump_eagle3_batch class Eagle3Trainer(Trainer): @@ -54,7 +54,7 @@ def init_model( mooncake_config=None, ) -> int: if mooncake_config is not None: - from torchspec.transfer.mooncake.utils import ( + from aurora.transfer.mooncake.utils import ( check_mooncake_master_available, ) @@ -155,7 +155,7 @@ def init_model( self.target_lm_head_weight = self.target_lm_head.lm_head.weight if getattr(self.args, "attention_backend", None) == "fa_experimental": - from torchspec.models.draft.llama3_eagle import ( + from aurora.models.draft.llama3_eagle import ( _has_cute_dsl, warmup_flash_attention_masked, ) @@ -188,7 +188,7 @@ def _init_target_lm_head(self, target_model_path: str) -> None: Only rank 0 loads the weights, then broadcasts to other ranks. The lm_head is kept frozen and not wrapped with FSDP. """ - from torchspec.models.target.target_utils import TargetLMHead + from aurora.models.target.target_utils import TargetLMHead if dist.get_rank() == 0: self.target_lm_head = TargetLMHead.from_pretrained( diff --git a/torchspec/training/fsdp.py b/aurora/training/fsdp.py similarity index 99% rename from torchspec/training/fsdp.py rename to aurora/training/fsdp.py index 0ba61f5..e3381b6 100644 --- a/torchspec/training/fsdp.py +++ b/aurora/training/fsdp.py @@ -25,7 +25,7 @@ import torch.distributed as dist import torch.nn as nn -from torchspec.utils.logging import logger +from aurora.utils.logging import logger @contextmanager diff --git a/torchspec/training/lr_scheduler.py b/aurora/training/lr_scheduler.py similarity index 98% rename from torchspec/training/lr_scheduler.py rename to aurora/training/lr_scheduler.py index 5a0221f..cf5f73f 100644 --- a/torchspec/training/lr_scheduler.py +++ b/aurora/training/lr_scheduler.py @@ -18,7 +18,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""Unified learning rate schedulers for torchspec training.""" +"""Unified learning rate schedulers for aurora training.""" import math from typing import Literal, Optional @@ -26,7 +26,7 @@ import torch from torch.optim.lr_scheduler import LRScheduler -from torchspec.utils.logging import logger +from aurora.utils.logging import logger DecayStyle = Literal["constant", "linear", "cosine", "inverse-square-root", "WSD"] WSDDecayStyle = Literal["linear", "cosine", "exponential", "minus_sqrt"] diff --git a/torchspec/training/optimizer.py b/aurora/training/optimizer.py similarity index 97% rename from torchspec/training/optimizer.py rename to aurora/training/optimizer.py index 5c1dc8b..eb4c95f 100644 --- a/torchspec/training/optimizer.py +++ b/aurora/training/optimizer.py @@ -20,8 +20,8 @@ import torch -from torchspec.training.lr_scheduler import LRSchedulerWithWarmup -from torchspec.utils.logging import print_on_rank0 +from aurora.training.lr_scheduler import LRSchedulerWithWarmup +from aurora.utils.logging import print_on_rank0 class BF16Optimizer: diff --git a/torchspec/training/trainer.py b/aurora/training/trainer.py similarity index 96% rename from torchspec/training/trainer.py rename to aurora/training/trainer.py index a77decf..a46a137 100644 --- a/torchspec/training/trainer.py +++ b/aurora/training/trainer.py @@ -35,17 +35,17 @@ ) from torch.distributed.device_mesh import init_device_mesh -from torchspec.config.mooncake_config import MooncakeConfig -from torchspec.data.utils import DataCollatorWithPadding -from torchspec.training import checkpoint -from torchspec.training.data_fetcher import MooncakeDataFetcher -from torchspec.training.fsdp import init_empty_weights -from torchspec.training.optimizer import BF16Optimizer -from torchspec.transfer.mooncake.eagle_store import EagleMooncakeStore -from torchspec.utils.logging import logger -from torchspec.utils.processing import get_assistant_token_ids -from torchspec.utils.profiling import TrainProfiler -from torchspec.utils.train_dump import extract_gradients, extract_model_weights +from aurora.config.mooncake_config import MooncakeConfig +from aurora.data.utils import DataCollatorWithPadding +from aurora.training import checkpoint +from aurora.training.data_fetcher import MooncakeDataFetcher +from aurora.training.fsdp import init_empty_weights +from aurora.training.optimizer import BF16Optimizer +from aurora.transfer.mooncake.eagle_store import EagleMooncakeStore +from aurora.utils.logging import logger +from aurora.utils.processing import get_assistant_token_ids +from aurora.utils.profiling import TrainProfiler +from aurora.utils.train_dump import extract_gradients, extract_model_weights class Trainer(abc.ABC): diff --git a/torchspec/training/trainer_actor.py b/aurora/training/trainer_actor.py similarity index 94% rename from torchspec/training/trainer_actor.py rename to aurora/training/trainer_actor.py index cfe114f..42261ca 100644 --- a/torchspec/training/trainer_actor.py +++ b/aurora/training/trainer_actor.py @@ -24,11 +24,11 @@ import torch.distributed as dist -from torchspec import AutoDraftModelConfig -from torchspec.ray.ray_actor import RayActor -from torchspec.training.eagle3_trainer import Eagle3Trainer -from torchspec.utils.distributed import init_gloo_group -from torchspec.utils.logging import setup_file_logging +from aurora import AutoDraftModelConfig +from aurora.ray.ray_actor import RayActor +from aurora.training.eagle3_trainer import Eagle3Trainer +from aurora.utils.distributed import init_gloo_group +from aurora.utils.logging import setup_file_logging class TrainerActor(RayActor): diff --git a/torchspec/transfer/__init__.py b/aurora/transfer/__init__.py similarity index 100% rename from torchspec/transfer/__init__.py rename to aurora/transfer/__init__.py diff --git a/torchspec/transfer/mooncake/__init__.py b/aurora/transfer/mooncake/__init__.py similarity index 84% rename from torchspec/transfer/mooncake/__init__.py rename to aurora/transfer/mooncake/__init__.py index ba0f62d..ec259a7 100644 --- a/torchspec/transfer/mooncake/__init__.py +++ b/aurora/transfer/mooncake/__init__.py @@ -18,8 +18,8 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -from torchspec.transfer.mooncake.helpers import calculate_eagle3_buffer_size -from torchspec.transfer.mooncake.utils import ( +from aurora.transfer.mooncake.helpers import calculate_eagle3_buffer_size +from aurora.transfer.mooncake.utils import ( MooncakeMaster, check_mooncake_master_available, launch_mooncake_master, @@ -30,15 +30,15 @@ def __getattr__(name): # Lazy imports to avoid circular dependency with config.mooncake_config if name == "MooncakeConfig": - from torchspec.config.mooncake_config import MooncakeConfig + from aurora.config.mooncake_config import MooncakeConfig return MooncakeConfig if name == "MooncakeHiddenStateStore": - from torchspec.transfer.mooncake.store import MooncakeHiddenStateStore + from aurora.transfer.mooncake.store import MooncakeHiddenStateStore return MooncakeHiddenStateStore if name == "EagleMooncakeStore": - from torchspec.transfer.mooncake.eagle_store import EagleMooncakeStore + from aurora.transfer.mooncake.eagle_store import EagleMooncakeStore return EagleMooncakeStore raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/torchspec/transfer/mooncake/buffers.py b/aurora/transfer/mooncake/buffers.py similarity index 99% rename from torchspec/transfer/mooncake/buffers.py rename to aurora/transfer/mooncake/buffers.py index ce7fadf..7a63991 100644 --- a/torchspec/transfer/mooncake/buffers.py +++ b/aurora/transfer/mooncake/buffers.py @@ -24,7 +24,7 @@ import torch -from torchspec.utils.logging import logger +from aurora.utils.logging import logger class HostBuffer: diff --git a/torchspec/transfer/mooncake/deferred_delete.py b/aurora/transfer/mooncake/deferred_delete.py similarity index 99% rename from torchspec/transfer/mooncake/deferred_delete.py rename to aurora/transfer/mooncake/deferred_delete.py index e037cb9..a552e95 100644 --- a/torchspec/transfer/mooncake/deferred_delete.py +++ b/aurora/transfer/mooncake/deferred_delete.py @@ -26,7 +26,7 @@ from dataclasses import dataclass from typing import Any, Dict, List -from torchspec.utils.logging import logger +from aurora.utils.logging import logger @dataclass diff --git a/torchspec/transfer/mooncake/eagle_store.py b/aurora/transfer/mooncake/eagle_store.py similarity index 98% rename from torchspec/transfer/mooncake/eagle_store.py rename to aurora/transfer/mooncake/eagle_store.py index 1dcd583..9ce3666 100644 --- a/torchspec/transfer/mooncake/eagle_store.py +++ b/aurora/transfer/mooncake/eagle_store.py @@ -25,13 +25,13 @@ import torch -from torchspec.transfer.mooncake.deferred_delete import DeferredDeleteManager -from torchspec.transfer.mooncake.helpers import _format_bytes -from torchspec.transfer.mooncake.store import MooncakeHiddenStateStore -from torchspec.utils.logging import logger +from aurora.transfer.mooncake.deferred_delete import DeferredDeleteManager +from aurora.transfer.mooncake.helpers import _format_bytes +from aurora.transfer.mooncake.store import MooncakeHiddenStateStore +from aurora.utils.logging import logger if TYPE_CHECKING: - from torchspec.models.target.eagle3_target_model import Eagle3TargetOutput + from aurora.models.target.eagle3_target_model import Eagle3TargetOutput # Static lookup for dtype → element size in bytes (avoids creating a tensor # on every call to _compute_tensor_size). @@ -273,7 +273,7 @@ def get( Returns: Eagle3TargetOutput with the retrieved tensors. """ - from torchspec.models.target.eagle3_target_model import Eagle3TargetOutput + from aurora.models.target.eagle3_target_model import Eagle3TargetOutput keys = [f"{key}_hs", f"{key}_ids"] tensor_specs = [ diff --git a/torchspec/transfer/mooncake/helpers.py b/aurora/transfer/mooncake/helpers.py similarity index 98% rename from torchspec/transfer/mooncake/helpers.py rename to aurora/transfer/mooncake/helpers.py index a614eab..c6e9e63 100644 --- a/torchspec/transfer/mooncake/helpers.py +++ b/aurora/transfer/mooncake/helpers.py @@ -18,7 +18,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -from torchspec.utils.logging import logger +from aurora.utils.logging import logger def _format_bytes(size: int) -> str: diff --git a/torchspec/transfer/mooncake/store.py b/aurora/transfer/mooncake/store.py similarity index 98% rename from torchspec/transfer/mooncake/store.py rename to aurora/transfer/mooncake/store.py index d416c52..54c0ccb 100644 --- a/torchspec/transfer/mooncake/store.py +++ b/aurora/transfer/mooncake/store.py @@ -24,14 +24,14 @@ import torch from mooncake.store import MooncakeDistributedStore -from torchspec.config.mooncake_config import MooncakeConfig -from torchspec.transfer.mooncake.buffers import ( +from aurora.config.mooncake_config import MooncakeConfig +from aurora.transfer.mooncake.buffers import ( AsyncPutManager, GPUReceiveBuffer, GPUSendBuffer, HostBufferPool, ) -from torchspec.utils.logging import logger +from aurora.utils.logging import logger class MooncakeHiddenStateStore(ABC): diff --git a/torchspec/transfer/mooncake/utils.py b/aurora/transfer/mooncake/utils.py similarity index 97% rename from torchspec/transfer/mooncake/utils.py rename to aurora/transfer/mooncake/utils.py index d052de5..d464ba3 100644 --- a/torchspec/transfer/mooncake/utils.py +++ b/aurora/transfer/mooncake/utils.py @@ -32,9 +32,9 @@ import ray -from torchspec.ray.ray_actor import RayActor -from torchspec.utils.env import get_torchspec_env_vars -from torchspec.utils.logging import logger +from aurora.ray.ray_actor import RayActor +from aurora.utils.env import get_aurora_env_vars +from aurora.utils.logging import logger def resolve_mooncake_master_bin() -> str: @@ -254,7 +254,7 @@ def launch_mooncake_master(args): Returns: The MooncakeMasterActor handle, or None if binary not found. """ - from torchspec.ray.ray_actor import node_affinity_for_ip + from aurora.ray.ray_actor import node_affinity_for_ip master_addr = getattr(args, "mooncake_master_server_address", None) scheduling_strategy = None @@ -302,7 +302,7 @@ def launch_mooncake_master(args): logger.warning(f"Binary not found at {mooncake_bin}, skipping launch") return None - RemoteActor = ray.remote(num_cpus=0, runtime_env={"env_vars": get_torchspec_env_vars()})( + RemoteActor = ray.remote(num_cpus=0, runtime_env={"env_vars": get_aurora_env_vars()})( MooncakeMaster ) actor_options = {"name": "mooncake_master"} @@ -310,7 +310,7 @@ def launch_mooncake_master(args): actor_options["scheduling_strategy"] = scheduling_strategy actor = RemoteActor.options(**actor_options).remote() - from torchspec.config.mooncake_config import MooncakeConfig + from aurora.config.mooncake_config import MooncakeConfig kv_lease_ttl_ms = getattr(args, "mooncake_kv_lease_ttl_ms", MooncakeConfig.kv_lease_ttl_ms) diff --git a/torchspec/utils/__init__.py b/aurora/utils/__init__.py similarity index 96% rename from torchspec/utils/__init__.py rename to aurora/utils/__init__.py index 17f232a..dfc51dd 100644 --- a/torchspec/utils/__init__.py +++ b/aurora/utils/__init__.py @@ -18,4 +18,4 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""Utility package for torchspec.""" +"""Utility package for aurora.""" diff --git a/torchspec/utils/distributed.py b/aurora/utils/distributed.py similarity index 100% rename from torchspec/utils/distributed.py rename to aurora/utils/distributed.py diff --git a/torchspec/utils/env.py b/aurora/utils/env.py similarity index 85% rename from torchspec/utils/env.py rename to aurora/utils/env.py index 6dad7f2..c1533d4 100644 --- a/torchspec/utils/env.py +++ b/aurora/utils/env.py @@ -4,7 +4,7 @@ # NOTE: TORCHINDUCTOR_CACHE_DIR is intentionally excluded — each node should # use its own node-local default (/tmp/torchinductor_$USER/) to avoid # cross-node triton kernel cache corruption over NFS. -_TORCHSPEC_ENV_KEYS = [ +_AURORA_ENV_KEYS = [ "CUDA_LAUNCH_BLOCKING", "GLOO_SOCKET_IFNAME", "HF_HOME", @@ -15,8 +15,8 @@ "NCCL_SOCKET_IFNAME", "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN", "SGLANG_DISABLE_CUDNN_CHECK", - "TORCHSPEC_LOG_DIR", - "TORCHSPEC_LOG_LEVEL", + "AURORA_LOG_DIR", + "AURORA_LOG_LEVEL", "TP_SOCKET_IFNAME", ] @@ -33,11 +33,11 @@ ] -def get_torchspec_env_vars() -> dict[str, str]: +def get_aurora_env_vars() -> dict[str, str]: """Return common environment variables for all Ray actors. Includes: - - TORCHSPEC_* variables (e.g. log level) from the current process + - AURORA_* variables (e.g. log level) from the current process - RAY_EXPERIMENTAL_NOSET_*_VISIBLE_DEVICES = "1" to prevent Ray from overriding device visibility @@ -45,5 +45,5 @@ def get_torchspec_env_vars() -> dict[str, str]: Call-site env vars merged after this dict take higher priority. """ env = {k: "1" for k in _RAY_NOSET_VISIBLE_DEVICES_KEYS} - env.update({k: os.environ[k] for k in _TORCHSPEC_ENV_KEYS if k in os.environ}) + env.update({k: os.environ[k] for k in _AURORA_ENV_KEYS if k in os.environ}) return env diff --git a/torchspec/utils/logging.py b/aurora/utils/logging.py similarity index 94% rename from torchspec/utils/logging.py rename to aurora/utils/logging.py index 79f227e..d90699a 100644 --- a/torchspec/utils/logging.py +++ b/aurora/utils/logging.py @@ -24,13 +24,13 @@ import torch.distributed as dist -from torchspec.utils import wandb as wandb_utils +from aurora.utils import wandb as wandb_utils _LOG_FORMAT = "[%(asctime)s] %(filename)s:%(lineno)d %(levelname)s %(message)s" def _get_logger_level(): - level_str = os.getenv("TORCHSPEC_LOG_LEVEL", "INFO").upper() + level_str = os.getenv("AURORA_LOG_LEVEL", "INFO").upper() try: log_level = getattr(logging, level_str) except ValueError: @@ -72,7 +72,7 @@ def setup_file_logging( ) -> None: """Add a FileHandler to the module-level logger for per-role/per-node/per-rank file logging.""" if log_dir is None: - log_dir = os.environ.get("TORCHSPEC_LOG_DIR") + log_dir = os.environ.get("AURORA_LOG_DIR") if log_dir is None: return @@ -82,7 +82,7 @@ def setup_file_logging( logger.removeHandler(h) try: - from torchspec.utils.misc import get_current_node_ip + from aurora.utils.misc import get_current_node_ip node_ip = get_current_node_ip() except Exception: diff --git a/torchspec/utils/memory.py b/aurora/utils/memory.py similarity index 98% rename from torchspec/utils/memory.py rename to aurora/utils/memory.py index 05f7331..d407c03 100644 --- a/torchspec/utils/memory.py +++ b/aurora/utils/memory.py @@ -24,7 +24,7 @@ import torch import torch.distributed as dist -from torchspec.utils.logging import logger +from aurora.utils.logging import logger DTYPE_SIZES = { torch.float32: 4, diff --git a/torchspec/utils/misc.py b/aurora/utils/misc.py similarity index 100% rename from torchspec/utils/misc.py rename to aurora/utils/misc.py diff --git a/torchspec/utils/processing.py b/aurora/utils/processing.py similarity index 95% rename from torchspec/utils/processing.py rename to aurora/utils/processing.py index 0ff8d39..b7ef011 100644 --- a/torchspec/utils/processing.py +++ b/aurora/utils/processing.py @@ -20,7 +20,7 @@ from transformers import AutoTokenizer -from torchspec.utils.logging import logger +from aurora.utils.logging import logger def load_tokenizer(name_or_path: str, **kwargs): @@ -30,7 +30,7 @@ def load_tokenizer(name_or_path: str, **kwargs): def get_assistant_token_ids(args) -> tuple[list[int] | None, list[int] | None]: """Derive assistant_header_ids and end_token_ids from chat_template config.""" - from torchspec.data.template import TEMPLATE_REGISTRY + from aurora.data.template import TEMPLATE_REGISTRY chat_template_name = getattr(args, "chat_template", None) if not chat_template_name: diff --git a/torchspec/utils/profiling.py b/aurora/utils/profiling.py similarity index 98% rename from torchspec/utils/profiling.py rename to aurora/utils/profiling.py index 1e19793..fd46b36 100644 --- a/torchspec/utils/profiling.py +++ b/aurora/utils/profiling.py @@ -24,8 +24,8 @@ import torch -from torchspec.utils.logging import logger -from torchspec.utils.memory import print_memory +from aurora.utils.logging import logger +from aurora.utils.memory import print_memory class TrainProfiler: diff --git a/torchspec/utils/tensor.py b/aurora/utils/tensor.py similarity index 100% rename from torchspec/utils/tensor.py rename to aurora/utils/tensor.py diff --git a/torchspec/utils/train_dump.py b/aurora/utils/train_dump.py similarity index 98% rename from torchspec/utils/train_dump.py rename to aurora/utils/train_dump.py index 1339543..b53a7cb 100644 --- a/torchspec/utils/train_dump.py +++ b/aurora/utils/train_dump.py @@ -29,7 +29,7 @@ import torch import torch.nn as nn -from torchspec.utils.logging import logger +from aurora.utils.logging import logger def extract_gradients(model: nn.Module) -> Dict[str, torch.Tensor]: @@ -410,21 +410,21 @@ def run_comparison_test( print("To compare training outputs from two configs:") print() print("1. Run config1 with debug dumps enabled:") - print(f" python -m torchspec.train.train_eagle3 --config {config1_path} \\") + print(f" python -m aurora.train.train_eagle3 --config {config1_path} \\") print( f" --save_debug_train_data '{dir1}/batch_{{step}}_{{batch_idx}}_rank{{rank}}.pt' \\" ) print(f" --max_num_steps {num_steps}") print() print("2. Run config2 with debug dumps enabled:") - print(f" python -m torchspec.train_entry --config {config2_path} \\") + print(f" python -m aurora.train_entry --config {config2_path} \\") print( f" debug.save_debug_train_data='{dir2}/batch_{{step}}_{{batch_idx}}_rank{{rank}}.pt' \\" ) print(f" training.num_train_steps={num_steps}") print() print("3. Compare results:") - print(" from torchspec.utils.train_dump import compare_eagle3_batches") + print(" from aurora.utils.train_dump import compare_eagle3_batches") print(f" results = compare_eagle3_batches('{dir1}', '{dir2}')") print() diff --git a/torchspec/utils/types.py b/aurora/utils/types.py similarity index 100% rename from torchspec/utils/types.py rename to aurora/utils/types.py diff --git a/torchspec/utils/wandb.py b/aurora/utils/wandb.py similarity index 100% rename from torchspec/utils/wandb.py rename to aurora/utils/wandb.py diff --git a/docs/code_architecture.md b/docs/code_architecture.md index d29afaf..33d2bc5 100644 --- a/docs/code_architecture.md +++ b/docs/code_architecture.md @@ -3,7 +3,7 @@ ## Package Layout ``` -torchspec/ +aurora/ ├── config/ # Configuration system (OmegaConf-based) │ ├── train_config.py # Hierarchical dataclass configs (7 sections + Config root) │ ├── inference_config.py # InferenceConfig + SGLangConfig (essential fields + extra_args passthrough) @@ -63,7 +63,7 @@ torchspec/ │ └── utils.py # Loss mask packing/unpacking ├── utils/ # Shared utilities │ ├── distributed.py # Device mesh setup, TP/DP primitives (get_tp_group, get_tp_device_mesh) -│ ├── env.py # Ray actor env-var forwarding (get_torchspec_env_vars) +│ ├── env.py # Ray actor env-var forwarding (get_aurora_env_vars) │ ├── logging.py # Unified logger │ ├── memory.py # Tensor byte estimation │ ├── profiling.py # PyTorch profiler utilities @@ -78,7 +78,7 @@ torchspec/ ## Core Components -### 1. Draft Model (`torchspec/models/draft/`) +### 1. Draft Model (`aurora/models/draft/`) A lightweight transformer initialized from the target model's architecture: @@ -90,7 +90,7 @@ A lightweight transformer initialized from the target model's architecture: - Hidden state projection from target model - Token-to-draft vocabulary mapping (`t2d`) -### 2. Target Model (`torchspec/models/target/`) +### 2. Target Model (`aurora/models/target/`) Abstract interface for running the target model during inference: @@ -102,7 +102,7 @@ The target model extracts: - **Hidden states** from configurable layers (`aux_hidden_states_layers`) - **Logits** for computing soft labels (KL divergence targets) -### 3. Async Training Controller (`torchspec/controller/training_controller.py`) +### 3. Async Training Controller (`aurora/controller/training_controller.py`) Central orchestrator (Ray actor) managing the async pipeline: @@ -123,13 +123,13 @@ class AsyncTrainingController: The controller only manages metadata and Mooncake keys, never actual tensor data. It tracks exact bytes in the sample pool for Mooncake backpressure control. -### 4. Async Inference Manager (`torchspec/controller/inference_manager.py`) +### 4. Async Inference Manager (`aurora/controller/inference_manager.py`) Self-regulating inference manager (Ray actor) that dispatches to `HFEngine` / `SglEngine` Ray actors with load balancing. Includes Mooncake backpressure: pauses generation when `sample_pool` exceeds capacity, resuming when training catches up. -### 5. Inference Engines (`torchspec/inference/engine/`) +### 5. Inference Engines (`aurora/inference/engine/`) - **`base.py`**: `InferenceEngine` - Abstract base class defining the unified engine interface - **`hf_runner.py`**: `HFRunner` - Core inference logic that runs target model, extracts hidden states, and stores tensors in Mooncake @@ -138,7 +138,7 @@ Includes Mooncake backpressure: pauses generation when `sample_pool` exceeds cap Factory function in `factory.py`: `create_inference_engines()` -### 6. Training (`torchspec/training/`) +### 6. Training (`aurora/training/`) The training side is split across three layers: @@ -147,7 +147,7 @@ The training side is split across three layers: - **`eagle3_trainer.py`**: `Eagle3Trainer(Trainer)` — Eagle3-specific logic: initialises `Eagle3Model` with the draft model under FSDP2, runs the forward/backward, and aggregates metrics - **`fsdp.py`**: FSDP2 helpers (`apply_fsdp2`, `fsdp2_load_full_state_dict`, `init_empty_weights`) -### 7. Mooncake Integration (`torchspec/transfer/mooncake/`) +### 7. Mooncake Integration (`aurora/transfer/mooncake/`) Distributed tensor transfer for multi-node training: @@ -187,7 +187,7 @@ Distributed tensor transfer for multi-node training: └── Periodic checkpointing ``` -## Configuration System (`torchspec/config/`) +## Configuration System (`aurora/config/`) Hierarchical YAML configs powered by OmegaConf, with 9 typed dataclass sections: @@ -225,7 +225,7 @@ mooncake: logging: report_to: wandb - wandb_project: torchspec + wandb_project: aurora debug: use_pytorch_profiler: false @@ -242,85 +242,85 @@ python train.py --config base.yaml --config experiment.yaml training.learning_ra | Module | Purpose | |--------|---------| -| `torchspec/models/eagle3.py` | `Eagle3Model` - Eagle3 forward pass and loss computation | -| `torchspec/models/ops/loss.py` | `compiled_forward_kl_loss` - Forward KL loss | -| `torchspec/models/ops/loss_mask.py` | Loss mask computation utilities | -| `torchspec/models/ops/flex_attention.py` | FlexAttention utilities | -| `torchspec/models/draft/auto.py` | `AutoEagle3DraftModel` factory | -| `torchspec/models/draft/base.py` | `Eagle3DraftModel` abstract base | -| `torchspec/models/draft/llama3_eagle.py` | `LlamaForCausalLMEagle3` implementation | -| `torchspec/models/target/eagle3_target_model.py` | `Eagle3TargetModel` ABC + `HFTargetModel` implementation | -| `torchspec/models/target/target_utils.py` | Hidden state layer selection utilities | +| `aurora/models/eagle3.py` | `Eagle3Model` - Eagle3 forward pass and loss computation | +| `aurora/models/ops/loss.py` | `compiled_forward_kl_loss` - Forward KL loss | +| `aurora/models/ops/loss_mask.py` | Loss mask computation utilities | +| `aurora/models/ops/flex_attention.py` | FlexAttention utilities | +| `aurora/models/draft/auto.py` | `AutoEagle3DraftModel` factory | +| `aurora/models/draft/base.py` | `Eagle3DraftModel` abstract base | +| `aurora/models/draft/llama3_eagle.py` | `LlamaForCausalLMEagle3` implementation | +| `aurora/models/target/eagle3_target_model.py` | `Eagle3TargetModel` ABC + `HFTargetModel` implementation | +| `aurora/models/target/target_utils.py` | Hidden state layer selection utilities | ### Ray Infrastructure | Module | Purpose | |--------|---------| -| `torchspec/ray/ray_actor.py` | `RayActor` base class (GPU setup, IP/port utils, master addr negotiation) | -| `torchspec/ray/train_group.py` | `RayTrainGroup` - Manages a group of training actors | -| `torchspec/ray/placement_group.py` | Placement group creation, GPU resource waiting, `create_placement_groups()`, `create_train_group()` | +| `aurora/ray/ray_actor.py` | `RayActor` base class (GPU setup, IP/port utils, master addr negotiation) | +| `aurora/ray/train_group.py` | `RayTrainGroup` - Manages a group of training actors | +| `aurora/ray/placement_group.py` | Placement group creation, GPU resource waiting, `create_placement_groups()`, `create_train_group()` | ### Controller | Module | Purpose | |--------|---------| -| `torchspec/controller/training_controller.py` | `AsyncTrainingController` - Pipeline orchestration | -| `torchspec/controller/inference_manager.py` | `AsyncInferenceManager` - Inference dispatch and backpressure | -| `torchspec/controller/loop.py` | `run_training_loop()` - Main training loop | -| `torchspec/controller/setup.py` | `build_mooncake_config`, `setup_async_training_with_engines`, `auto_calculate_training_steps` | +| `aurora/controller/training_controller.py` | `AsyncTrainingController` - Pipeline orchestration | +| `aurora/controller/inference_manager.py` | `AsyncInferenceManager` - Inference dispatch and backpressure | +| `aurora/controller/loop.py` | `run_training_loop()` - Main training loop | +| `aurora/controller/setup.py` | `build_mooncake_config`, `setup_async_training_with_engines`, `auto_calculate_training_steps` | ### Inference | Module | Purpose | |--------|---------| -| `torchspec/inference/factory.py` | `create_inference_engines()` - Engine creation with placement groups | -| `torchspec/inference/engine/base.py` | `InferenceEngine` abstract base class | -| `torchspec/inference/engine/hf_runner.py` | `HFRunner` core inference logic | -| `torchspec/inference/engine/hf_engine.py` | `HFEngine` Ray actor wrapper (inherits `RayActor`) | -| `torchspec/inference/engine/sgl_engine.py` | `SglEngine` Ray actor wrapper (inherits `RayActor`) | +| `aurora/inference/factory.py` | `create_inference_engines()` - Engine creation with placement groups | +| `aurora/inference/engine/base.py` | `InferenceEngine` abstract base class | +| `aurora/inference/engine/hf_runner.py` | `HFRunner` core inference logic | +| `aurora/inference/engine/hf_engine.py` | `HFEngine` Ray actor wrapper (inherits `RayActor`) | +| `aurora/inference/engine/sgl_engine.py` | `SglEngine` Ray actor wrapper (inherits `RayActor`) | ### Training | Module | Purpose | |--------|-------| -| `torchspec/training/trainer_actor.py` | `TrainerActor` - Ray actor wrapper; owns distributed process group | -| `torchspec/training/trainer.py` | `Trainer` - Abstract base (device mesh, data fetcher, loop skeleton) | -| `torchspec/training/eagle3_trainer.py` | `Eagle3Trainer` - Eagle3 model init, forward/backward, metric aggregation | -| `torchspec/training/fsdp.py` | `apply_fsdp2`, `fsdp2_load_full_state_dict`, `init_empty_weights` | -| `torchspec/training/data_fetcher.py` | `MooncakeDataFetcher` - Queue-based data retrieval | -| `torchspec/training/checkpoint.py` | Checkpoint save/load | -| `torchspec/training/optimizer.py` | `BF16Optimizer` construction | -| `torchspec/training/lr_scheduler.py` | LR scheduling | +| `aurora/training/trainer_actor.py` | `TrainerActor` - Ray actor wrapper; owns distributed process group | +| `aurora/training/trainer.py` | `Trainer` - Abstract base (device mesh, data fetcher, loop skeleton) | +| `aurora/training/eagle3_trainer.py` | `Eagle3Trainer` - Eagle3 model init, forward/backward, metric aggregation | +| `aurora/training/fsdp.py` | `apply_fsdp2`, `fsdp2_load_full_state_dict`, `init_empty_weights` | +| `aurora/training/data_fetcher.py` | `MooncakeDataFetcher` - Queue-based data retrieval | +| `aurora/training/checkpoint.py` | Checkpoint save/load | +| `aurora/training/optimizer.py` | `BF16Optimizer` construction | +| `aurora/training/lr_scheduler.py` | LR scheduling | ### Data Pipeline | Module | Purpose | |--------|---------| -| `torchspec/data/dataset.py` | `load_conversation_dataset()` with format detection | -| `torchspec/data/parse.py` | Chat format parsers (`GeneralParser`, etc.) | -| `torchspec/data/preprocessing.py` | Tokenization, chat templates, loss masks | -| `torchspec/data/template.py` | Chat template handling | -| `torchspec/data/utils.py` | Loss mask packing/unpacking | +| `aurora/data/dataset.py` | `load_conversation_dataset()` with format detection | +| `aurora/data/parse.py` | Chat format parsers (`GeneralParser`, etc.) | +| `aurora/data/preprocessing.py` | Tokenization, chat templates, loss masks | +| `aurora/data/template.py` | Chat template handling | +| `aurora/data/utils.py` | Loss mask packing/unpacking | ### Configuration | Module | Purpose | |--------|-------| -| `torchspec/config/train_config.py` | `Config` root + 7 typed dataclass sections (`DatasetConfig`, `DebugConfig`, `InferenceConfig`, `LoggingConfig`, `ModelConfig`, `TrainingConfig`, plus `mooncake: dict`) | -| `torchspec/config/inference_config.py` | `InferenceConfig`, `SGLangConfig` (essential fields + `extra_args` passthrough), `HFInferenceConfig` | -| `torchspec/config/mooncake_config.py` | `MooncakeConfig` with env-var support and `from_flat_args()` | -| `torchspec/config/utils.py` | Config loading helpers, `generate_draft_model_config` | +| `aurora/config/train_config.py` | `Config` root + 7 typed dataclass sections (`DatasetConfig`, `DebugConfig`, `InferenceConfig`, `LoggingConfig`, `ModelConfig`, `TrainingConfig`, plus `mooncake: dict`) | +| `aurora/config/inference_config.py` | `InferenceConfig`, `SGLangConfig` (essential fields + `extra_args` passthrough), `HFInferenceConfig` | +| `aurora/config/mooncake_config.py` | `MooncakeConfig` with env-var support and `from_flat_args()` | +| `aurora/config/utils.py` | Config loading helpers, `generate_draft_model_config` | ### Infrastructure | Module | Purpose | |--------|-------| -| `torchspec/transfer/mooncake/` | Mooncake tensor transfer (RDMA/TCP, buffer pools, deferred delete) | -| `torchspec/utils/distributed.py` | Device mesh setup, TP/DP primitives (`get_tp_group`, `get_tp_device_mesh`) | -| `torchspec/utils/env.py` | Ray actor env-var forwarding (`get_torchspec_env_vars`) | -| `torchspec/utils/logging.py` | Unified logger | -| `torchspec/utils/profiling.py` | PyTorch profiler utilities | -| `torchspec/utils/types.py` | `InferenceInput`, `InferenceOutput` | -| `torchspec/utils/memory.py` | Tensor byte estimation | -| `torchspec/utils/wandb.py` | Weights & Biases integration | -| `torchspec/train_entry.py` | Main entry point (config parsing, Ray setup, launch) | +| `aurora/transfer/mooncake/` | Mooncake tensor transfer (RDMA/TCP, buffer pools, deferred delete) | +| `aurora/utils/distributed.py` | Device mesh setup, TP/DP primitives (`get_tp_group`, `get_tp_device_mesh`) | +| `aurora/utils/env.py` | Ray actor env-var forwarding (`get_aurora_env_vars`) | +| `aurora/utils/logging.py` | Unified logger | +| `aurora/utils/profiling.py` | PyTorch profiler utilities | +| `aurora/utils/types.py` | `InferenceInput`, `InferenceOutput` | +| `aurora/utils/memory.py` | Tensor byte estimation | +| `aurora/utils/wandb.py` | Weights & Biases integration | +| `aurora/train_entry.py` | Main entry point (config parsing, Ray setup, launch) | diff --git a/docs/debugging_ray_jobs.md b/docs/debugging_ray_jobs.md index 2bd141b..45f3231 100644 --- a/docs/debugging_ray_jobs.md +++ b/docs/debugging_ray_jobs.md @@ -64,13 +64,13 @@ tail "$RAY_TEMP_DIR/session_latest/logs/monitor.log" ### 4. Environment variable for log verbosity ```bash -export TORCHSPEC_LOG_LEVEL=DEBUG # default is INFO; DEBUG gives per-step detail +export AURORA_LOG_LEVEL=DEBUG # default is INFO; DEBUG gives per-step detail export NCCL_DEBUG=INFO # NCCL connection/transport debugging ``` ## Startup Sequence & What to Expect -The training entry point (`torchspec/train_entry.py → train_async_no_generation`) proceeds through these phases. If the job appears stuck, identify which phase it's in. +The training entry point (`aurora/train_entry.py → train_async_no_generation`) proceeds through these phases. If the job appears stuck, identify which phase it's in. ### Online training (internal SglEngine) diff --git a/docs/ray.md b/docs/ray.md index 1f1d6d6..ab8c2b0 100644 --- a/docs/ray.md +++ b/docs/ray.md @@ -7,24 +7,24 @@ Aurora uses Ray as its distributed orchestration layer. In **online training** m `RayActor` is the base class for all GPU-bound actors. It provides GPU setup, IP discovery, and port allocation so each actor doesn't reinvent them. ``` -torchspec/ray/ +aurora/ray/ ├── ray_actor.py RayActor base class ├── train_group.py RayTrainGroup (training actor group manager) └── placement_group.py Placement group creation & GPU resource management -torchspec/inference/engine/ +aurora/inference/engine/ ├── hf_engine.py HFEngine(InferenceEngine, RayActor) └── sgl_engine.py SglEngine(InferenceEngine, RayActor) -torchspec/training/ +aurora/training/ ├── trainer.py Trainer (ABC base) ├── trainer_actor.py TrainerActor(RayActor) — wraps Eagle3Trainer └── eagle3_trainer.py Eagle3Trainer(Trainer) — FSDP2 training logic -torchspec/transfer/mooncake/ +aurora/transfer/mooncake/ └── utils.py MooncakeMaster(RayActor) -torchspec/controller/ +aurora/controller/ ├── training_controller.py AsyncTrainingController (standalone Ray actor) ├── inference_manager.py AsyncInferenceManager (standalone Ray actor) └── training_external_server.py TrainingExternalServer (Ray actor — HTTP callback server for external sglang) @@ -96,11 +96,11 @@ ray start \ > **Important:** The example `run.sh` scripts are designed for single-node use — > they run `ray stop --force` and start their own local head node. For multi-node, -> invoke `torchspec.train_entry` directly against the pre-existing cluster: +> invoke `aurora.train_entry` directly against the pre-existing cluster: ```bash export RAY_ADDRESS=:6379 -python3 -m torchspec.train_entry --config [overrides...] +python3 -m aurora.train_entry --config [overrides...] ``` Aurora auto-detects the cluster via `RAY_ADDRESS`. Worker nodes don't @@ -117,7 +117,7 @@ manage their own local cluster): ```bash export RAY_ADDRESS=ray://:10001 -python3 -m torchspec.train_entry --config [overrides...] +python3 -m aurora.train_entry --config [overrides...] ``` ### NCCL / Gloo networking diff --git a/environment.yml b/environment.yml index 6d41901..c0533ba 100644 --- a/environment.yml +++ b/environment.yml @@ -1,4 +1,4 @@ -name: torchspec +name: aurora channels: - conda-forge dependencies: diff --git a/examples/README.md b/examples/README.md index 9fee8d5..8d9df13 100644 --- a/examples/README.md +++ b/examples/README.md @@ -54,64 +54,32 @@ Each example directory contains: ## Running an Example -### Single-node examples (most common) - -**Applies to:** all `*-external-no-draft` and `*-external-with-draft` folders except the `*-2node` variant. - -**Step 1** — Launch training + sglang (stays in foreground): - ```bash -bash examples//run.sh -``` +# 1. Start training + SGLang server +bash examples/qwen3-4b-external-no-draft/run.sh -`run.sh` orchestrates everything in order: Ray cluster → training → mooncake → callback server → (draft model creation, if applicable) → sglang server. It stays running and waits for training to finish. - -**Step 2** — In a **separate terminal**, send traffic once sglang is healthy: - -```bash -bash examples//send_requests.sh +# 2. In another terminal, send requests +bash examples/qwen3-4b-external-no-draft/send_requests.sh ``` -Only run this after `run.sh` prints that the sglang server is healthy. - -### 2-node examples - -**Applies to:** `qwen3-8b-coder-next-external-with-draft-2node`, `kimi-k25-nvfp4-external-no-draft`. - -Both nodes must share a filesystem (e.g., NFS) for the draft model checkpoint and weight sync. - -**Step 1** — On **Node 1** (trainer machine), launch training: - -```bash -NODE2_IP= bash examples//run_node1_train.sh # or run_trainer.sh -``` +### Multi-node (2-node) example -Wait until it prints **"Training callback server is ready"**. +The `*-2node` examples split training and inference across two machines. Both nodes must share a filesystem (e.g., NFS) for the draft model checkpoint and weight sync. -**Step 2** — On **Node 2** (inference machine), launch sglang: - -```bash -NODE1_IP= bash examples//run_node2_sglang.sh # or run_sglang.sh -``` - -This connects back to Node 1's mooncake/callback server. Wait until sglang is healthy. - -**Step 3** — Send traffic (from either node): - -```bash -SGLANG_URL=http://:30000 \ - bash examples//send_requests.sh -``` - -### Online training (no external sglang) - -**Applies to:** `qwen3-coder-next-online`. - -```bash -bash examples/qwen3-coder-next-online/run.sh -``` +A dataset must be provided to the trainer so it can build the vocab mapping before the draft model is created. The SGLang server on Node 2 runs independently — if the trainer on Node 1 crashes, the SGLang server continues serving requests unaffected. -That's it — inference runs embedded inside the training process via `SglEngine`, so there is no separate sglang server and no `send_requests.sh`. +1. **Machine 1:** Run `run_node1_train.sh` + ```bash + NODE2_IP= bash examples/qwen3-8b-coder-next-external-with-draft-2node/run_node1_train.sh + ``` +2. **Machine 2:** Run `run_node2_sglang.sh` (after Node 1 prints "Node 1 is ready") + ```bash + NODE1_IP= bash examples/qwen3-8b-coder-next-external-with-draft-2node/run_node2_sglang.sh + ``` +3. **Machine 2:** Run `send_requests.sh` + ```bash + SGLANG_URL=http://:30000 bash examples/qwen3-8b-coder-next-external-with-draft-2node/send_requests.sh + ``` ### Config overrides diff --git a/examples/kimi-k25-nvfp4-external-no-draft/run_trainer.sh b/examples/kimi-k25-nvfp4-external-no-draft/run_trainer.sh index deba2e1..2e30cbf 100755 --- a/examples/kimi-k25-nvfp4-external-no-draft/run_trainer.sh +++ b/examples/kimi-k25-nvfp4-external-no-draft/run_trainer.sh @@ -47,7 +47,7 @@ SGLANG_PORT="${SGLANG_PORT:-30000}" MOONCAKE_GRPC_PORT="${MOONCAKE_GRPC_PORT:-50052}" MOONCAKE_META_PORT="${MOONCAKE_META_PORT:-8090}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -96,7 +96,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$ROOT_DIR/outputs/kimi-k25-nvfp4-external-no-draft" \ diff --git a/examples/minimax-m21-external-no-draft/run.sh b/examples/minimax-m21-external-no-draft/run.sh index 26d2fcb..5eb663b 100755 --- a/examples/minimax-m21-external-no-draft/run.sh +++ b/examples/minimax-m21-external-no-draft/run.sh @@ -3,7 +3,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Starts a standalone sglang server WITHOUT speculative decoding # 5. Waits for the sglang server to be healthy @@ -56,7 +56,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -113,7 +113,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data_shuffled.jsonl" \ training.training_num_gpus_per_node="$TRAIN_GPUS" \ diff --git a/examples/minimax-m21-external-with-draft/run.sh b/examples/minimax-m21-external-with-draft/run.sh index b5a8a99..33d3574 100755 --- a/examples/minimax-m21-external-with-draft/run.sh +++ b/examples/minimax-m21-external-with-draft/run.sh @@ -3,7 +3,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Waits for auto-created scratch draft model # 5. Starts a standalone sglang server with EAGLE3 speculative decoding @@ -67,7 +67,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -124,7 +124,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data_shuffled.jsonl" \ output_dir="$OUTPUT_DIR" \ diff --git a/examples/qwen3-4b-external-no-draft/run.sh b/examples/qwen3-4b-external-no-draft/run.sh index 653e427..775e265 100755 --- a/examples/qwen3-4b-external-no-draft/run.sh +++ b/examples/qwen3-4b-external-no-draft/run.sh @@ -3,7 +3,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Starts a standalone sglang server WITHOUT speculative decoding # 5. Waits for the sglang server to be healthy @@ -56,7 +56,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -115,7 +115,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$ROOT_DIR/outputs/qwen3-4b-external-no-draft" \ diff --git a/examples/qwen3-4b-external-with-draft/run.sh b/examples/qwen3-4b-external-with-draft/run.sh index 5b3a7b2..04337f7 100755 --- a/examples/qwen3-4b-external-with-draft/run.sh +++ b/examples/qwen3-4b-external-with-draft/run.sh @@ -3,7 +3,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Waits for scratch draft model to be created by training # 5. Starts a standalone sglang server with EAGLE3 speculative decoding @@ -66,7 +66,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -125,7 +125,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$OUTPUT_DIR" \ diff --git a/examples/qwen3-8b-coder-next-external-no-draft/run.sh b/examples/qwen3-8b-coder-next-external-no-draft/run.sh index d562430..b42cebc 100755 --- a/examples/qwen3-8b-coder-next-external-no-draft/run.sh +++ b/examples/qwen3-8b-coder-next-external-no-draft/run.sh @@ -5,7 +5,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Starts a standalone sglang server WITHOUT speculative decoding # 5. Waits for the sglang server to be healthy @@ -57,7 +57,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -114,7 +114,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ training.training_num_gpus_per_node="$TRAIN_GPUS" \ diff --git a/examples/qwen3-8b-coder-next-external-with-draft-2node/config.yaml b/examples/qwen3-8b-coder-next-external-with-draft-2node/config.yaml index 1b06b07..94696fb 100644 --- a/examples/qwen3-8b-coder-next-external-with-draft-2node/config.yaml +++ b/examples/qwen3-8b-coder-next-external-with-draft-2node/config.yaml @@ -3,7 +3,7 @@ # Target model: Qwen/Qwen3-Coder-Next (80B MoE, 3B active params) # # Two-node setup: -# Node 1 (training): runs torchspec training actors + mooncake master + callback server +# Node 1 (training): runs aurora training actors + mooncake master + callback server # Node 2 (inference): runs standalone sglang server with EAGLE3 speculative decoding # # Requirements: diff --git a/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node1_train.sh b/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node1_train.sh index 1810dd2..567206b 100755 --- a/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node1_train.sh +++ b/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node1_train.sh @@ -6,7 +6,7 @@ # This script runs on the TRAINING node and: # 1. Starts Ray cluster (head node) # 2. Starts mooncake master -# 3. Starts torchspec training actors +# 3. Starts aurora training actors # 4. Starts the callback HTTP server (receives samples from Node 2's sglang) # 5. Periodically syncs draft model weights to Node 2 via shared filesystem # @@ -22,7 +22,7 @@ # # Environment variables: # NODE2_IP - (required) IP address of the inference node -# SHARED_DIR - Shared filesystem path for weight sync (default: /scratch/shared/torchspec) +# SHARED_DIR - Shared filesystem path for weight sync (default: /scratch/shared/aurora) # SGLANG_PORT - sglang server port on Node 2 (default: 30000) # CALLBACK_PORT - Training callback server port (default: 18080) @@ -58,12 +58,12 @@ MOONCAKE_GRPC_PORT="${MOONCAKE_GRPC_PORT:-50052}" MOONCAKE_META_PORT="${MOONCAKE_META_PORT:-8090}" # Output dir on shared filesystem (NFS) — must be accessible from both nodes # so Node 2 can read the scratch draft model and weight sync checkpoints. -OUTPUT_DIR="${OUTPUT_DIR:-/data/bobbie/tmp/torchspec/qwen3-next-coder-external-2node}" +OUTPUT_DIR="${OUTPUT_DIR:-/data/bobbie/tmp/aurora/qwen3-next-coder-external-2node}" SCRATCH_DRAFT_DIR="$OUTPUT_DIR/scratch_draft_model" -WEIGHT_SYNC_DIR="~/weight_sync" +WEIGHT_SYNC_DIR="$HOME/weight_sync" mkdir -p "$OUTPUT_DIR" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -116,7 +116,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # The sglang server on Node 2 connects to this node's mooncake master # and sends callbacks to this node's callback server. echo "Starting training..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$OUTPUT_DIR" \ diff --git a/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node2_sglang.sh b/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node2_sglang.sh index b87ffc1..76a3c6b 100755 --- a/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node2_sglang.sh +++ b/examples/qwen3-8b-coder-next-external-with-draft-2node/run_node2_sglang.sh @@ -21,7 +21,7 @@ # # Environment variables: # NODE1_IP - (required) IP address of the training node -# DRAFT_MODEL - Path to draft model (default: /scratch/shared/torchspec/scratch_draft_model) +# DRAFT_MODEL - Path to draft model (default: /scratch/shared/aurora/scratch_draft_model) # SGLANG_GPUS - GPUs for sglang server (default: 0,1,2,3) # SGLANG_PORT - sglang server port (default: 30000) # CALLBACK_PORT - Training callback port on Node 1 (default: 18080) @@ -49,7 +49,7 @@ MOONCAKE_META_PORT="${MOONCAKE_META_PORT:-8090}" TARGET_MODEL="${TARGET_MODEL:-/scratch/bobbie/hf_cache/Qwen3-Coder-Next}" # Shared filesystem (NFS) — must match OUTPUT_DIR from Node 1 -OUTPUT_DIR="${OUTPUT_DIR:-/data/bobbie/tmp/torchspec/qwen3-next-coder-external-2node}" +OUTPUT_DIR="${OUTPUT_DIR:-/data/bobbie/tmp/aurora/qwen3-next-coder-external-2node}" DRAFT_MODEL="${DRAFT_MODEL:-$OUTPUT_DIR/scratch_draft_model}" # Speculative decoding settings diff --git a/examples/qwen3-8b-coder-next-external-with-draft/run.sh b/examples/qwen3-8b-coder-next-external-with-draft/run.sh index 8269c50..743f240 100755 --- a/examples/qwen3-8b-coder-next-external-with-draft/run.sh +++ b/examples/qwen3-8b-coder-next-external-with-draft/run.sh @@ -5,7 +5,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Waits for auto-created scratch draft model # 5. Starts a standalone sglang server with EAGLE3 speculative decoding @@ -69,7 +69,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -126,7 +126,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$OUTPUT_DIR" \ diff --git a/examples/qwen3-8b-external-no-draft/run.sh b/examples/qwen3-8b-external-no-draft/run.sh index d36d825..c2b9c07 100755 --- a/examples/qwen3-8b-external-no-draft/run.sh +++ b/examples/qwen3-8b-external-no-draft/run.sh @@ -3,7 +3,7 @@ # # This script: # 1. Starts Ray cluster -# 2. Starts torchspec training (mooncake + callback server) +# 2. Starts aurora training (mooncake + callback server) # 3. Waits for the training callback server to be ready # 4. Starts a standalone sglang server WITHOUT speculative decoding # 5. Waits for the sglang server to be healthy @@ -56,7 +56,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -115,7 +115,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$ROOT_DIR/outputs/qwen3-8b-external-no-draft" \ diff --git a/examples/qwen3-8b-external-with-draft/config.yaml b/examples/qwen3-8b-external-with-draft/config.yaml index ec8225d..bb7817f 100644 --- a/examples/qwen3-8b-external-with-draft/config.yaml +++ b/examples/qwen3-8b-external-with-draft/config.yaml @@ -7,7 +7,7 @@ # - 2 GPUs for training (DP/FSDP: draft model sharded) # # Usage: -# python -m torchspec.train_entry --config configs/sglang_qwen3_8b_external.yaml +# python -m aurora.train_entry --config configs/sglang_qwen3_8b_external.yaml model: target_model_path: Qwen/Qwen3-8B diff --git a/examples/qwen3-8b-external-with-draft/run.sh b/examples/qwen3-8b-external-with-draft/run.sh index b477dc9..1a36601 100755 --- a/examples/qwen3-8b-external-with-draft/run.sh +++ b/examples/qwen3-8b-external-with-draft/run.sh @@ -2,7 +2,7 @@ # Train with SglEngine in decode mode + external sglang server (hybrid mode) # # This script: -# 1. Starts torchspec training (which starts mooncake + TrainingExternalServer) +# 1. Starts aurora training (which starts mooncake + TrainingExternalServer) # 2. Waits for the training callback server to be ready # 3. Starts a standalone sglang server with --spec-training-callback-url # 4. Waits for the sglang server to be healthy @@ -80,7 +80,7 @@ MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-12}" IFS=',' read -ra SGLANG_GPU_ARRAY <<< "$SGLANG_GPUS" SGLANG_TP_SIZE="${SGLANG_TP_SIZE:-${#SGLANG_GPU_ARRAY[@]}}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -140,7 +140,7 @@ ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-sta # --- Step 2: Start training in background (starts mooncake + callback server) --- echo "Starting training (mooncake master + callback server will come up)..." -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_train_data.jsonl" \ output_dir="$ROOT_DIR/outputs/qwen3-8b-external-with-draft" \ diff --git a/examples/qwen3-coder-next-online/run.sh b/examples/qwen3-coder-next-online/run.sh index ee3a730..8f348ea 100755 --- a/examples/qwen3-coder-next-online/run.sh +++ b/examples/qwen3-coder-next-online/run.sh @@ -31,7 +31,7 @@ TOTAL_GPUS=${#GPU_ARRAY[@]} TRAIN_GPUS="${TRAIN_GPUS:-2}" INFERENCE_GPUS="${INFERENCE_GPUS:-4}" -export TORCHSPEC_LOG_LEVEL=INFO +export AURORA_LOG_LEVEL=INFO LOG_DIR="$ROOT_DIR/running_logs" mkdir -p "$LOG_DIR" @@ -59,7 +59,7 @@ ray stop --force 2>/dev/null || true echo "Starting Ray on port $RAY_PORT with $TOTAL_GPUS GPUs..." ray start --head --num-gpus "$TOTAL_GPUS" --port "$RAY_PORT" --disable-usage-stats -python3 -m torchspec.train_entry \ +python3 -m aurora.train_entry \ --config "$CONFIG_FILE" \ dataset.train_data_path="$ROOT_DIR/datasets/onlinesd/merged/merged_code_train_shuffled.jsonl" \ training.training_num_gpus_per_node="$TRAIN_GPUS" \ diff --git a/examples/send_user_requests.py b/examples/send_user_requests.py index b79a265..10e895f 100755 --- a/examples/send_user_requests.py +++ b/examples/send_user_requests.py @@ -68,7 +68,7 @@ def load_dataset(path: str, prompt_key: str = "conversations"): def _normalize_conversation(conversation): """Normalize ShareGPT format (from/value) to standard (role/content). - Mirrors torchspec.data.preprocessing._normalize_conversation.""" + Mirrors aurora.data.preprocessing._normalize_conversation.""" ROLE_MAPPING = {"human": "user", "gpt": "assistant"} if not conversation: return conversation @@ -96,15 +96,15 @@ def _strip_trailing_assistant(messages): def format_prompt(messages, tokenizer, chat_template_name="qwen"): - """Apply chat template matching torchspec's GeneralParser.format(). + """Apply chat template matching aurora's GeneralParser.format(). - Steps (mirroring torchspec/data/parse.py GeneralParser.format): + Steps (mirroring aurora/data/parse.py GeneralParser.format): 1. Normalize ShareGPT format 2. Inject system prompt from template if not present 3. Strip trailing assistant messages (we want the model to generate) 4. Apply tokenizer chat template with add_generation_prompt=True """ - from torchspec.data.template import TEMPLATE_REGISTRY + from aurora.data.template import TEMPLATE_REGISTRY messages = _normalize_conversation(messages) messages = _strip_trailing_assistant(messages) diff --git a/patches/sglang/v0.5.8.post1/sglang.patch b/patches/sglang/v0.5.8.post1/sglang.patch index fd43e3a..50be1e6 100644 --- a/patches/sglang/v0.5.8.post1/sglang.patch +++ b/patches/sglang/v0.5.8.post1/sglang.patch @@ -690,7 +690,7 @@ index e818deaa4..4e2ac4a1c 100644 + self._decode_lhs_storage = {} + if self.server_args.enable_spec_training_mooncake: + try: -+ from torchspec.transfer.mooncake import ( ++ from aurora.transfer.mooncake import ( + EagleMooncakeStore, + MooncakeConfig, + ) @@ -702,7 +702,7 @@ index e818deaa4..4e2ac4a1c 100644 + logger.info("EagleMooncakeStore initialized for spec training") + except ImportError: + logger.warning( -+ "torchspec.mooncake not found. Spec training mooncake store disabled." ++ "aurora.mooncake not found. Spec training mooncake store disabled." + ) + def init_chunked_prefill(self): diff --git a/pyproject.toml b/pyproject.toml index 12a7721..6d0f431 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=61.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "torchspec" +name = "aurora" dynamic = ["version", "description"] readme = "README.md" requires-python = ">=3.12" @@ -32,7 +32,7 @@ dependencies = [ ] [tool.setuptools] -packages = ["torchspec"] +packages = ["aurora"] [project.optional-dependencies] dev = [ @@ -64,5 +64,5 @@ ignore = ["E501"] ban-relative-imports = "all" [tool.ruff.lint.isort] -known-first-party = ["torchspec"] +known-first-party = ["aurora"] known-third-party = ["wandb"] diff --git a/tests/test_capacity_handling.py b/tests/test_capacity_handling.py index 3bf4181..c2f148a 100644 --- a/tests/test_capacity_handling.py +++ b/tests/test_capacity_handling.py @@ -13,7 +13,7 @@ import pytest import torch -from torchspec.utils.memory import estimate_tensor_bytes +from aurora.utils.memory import estimate_tensor_bytes class TestEstimateTensorBytes: @@ -123,7 +123,7 @@ def _create_mock_inference_output( tensor_dtypes: dict, ): """Create a mock InferenceOutput.""" - from torchspec.utils.types import InferenceOutput + from aurora.utils.types import InferenceOutput return InferenceOutput( data_id=data_id, @@ -138,7 +138,7 @@ def _create_controller_class(): import importlib import sys - module_name = "torchspec.controller.training_controller" + module_name = "aurora.controller.training_controller" if module_name in sys.modules: del sys.modules[module_name] @@ -463,7 +463,7 @@ def test_push_result_with_empty_shapes(self): assert controller.get_pool_size() == 1 def test_push_result_with_none_shapes(self): - from torchspec.utils.types import InferenceOutput + from aurora.utils.types import InferenceOutput controller = self._create_controller() diff --git a/tests/test_data_fetcher.py b/tests/test_data_fetcher.py index 5e1b7c0..6f81054 100644 --- a/tests/test_data_fetcher.py +++ b/tests/test_data_fetcher.py @@ -6,7 +6,7 @@ import torch -from torchspec.training.data_fetcher import ( +from aurora.training.data_fetcher import ( MooncakeDataFetcher, MooncakeDataset, TrainSample, diff --git a/tests/test_hf_engine.py b/tests/test_hf_engine.py index 2365b4d..0494957 100644 --- a/tests/test_hf_engine.py +++ b/tests/test_hf_engine.py @@ -21,7 +21,7 @@ class MockArgs: def _import_hf_engine(): """Import HFEngine, skipping test if dependencies unavailable.""" try: - from torchspec.inference.engine.hf_engine import HFEngine + from aurora.inference.engine.hf_engine import HFEngine return HFEngine except ImportError as e: @@ -31,7 +31,7 @@ def _import_hf_engine(): def _get_engine_module(): """Get the hf_engine module, skipping test if unavailable.""" try: - import torchspec.inference.engine.hf_engine as engine_module + import aurora.inference.engine.hf_engine as engine_module return engine_module except ImportError as e: @@ -168,10 +168,10 @@ def test_init_creates_engine_without_mooncake(self): mock_engine_class.from_pretrained.return_value = mock_inference_engine with patch( - "torchspec.inference.engine.hf_runner.HFRunner", + "aurora.inference.engine.hf_runner.HFRunner", mock_engine_class, ): - with patch("torchspec.config.mooncake_config.MooncakeConfig", MagicMock()): + with patch("aurora.config.mooncake_config.MooncakeConfig", MagicMock()): engine.init(mooncake_config=None) assert engine._engine is mock_inference_engine @@ -193,10 +193,10 @@ def test_init_creates_engine_with_mooncake(self): mooncake_dict = {"master_server_address": "localhost:50051"} with patch( - "torchspec.inference.engine.hf_runner.HFRunner", + "aurora.inference.engine.hf_runner.HFRunner", mock_engine_class, ): - with patch("torchspec.config.mooncake_config.MooncakeConfig", mock_mooncake_config): + with patch("aurora.config.mooncake_config.MooncakeConfig", mock_mooncake_config): engine.init(mooncake_config=mooncake_dict) assert engine._engine is mock_inference_engine @@ -214,12 +214,12 @@ def test_init_sets_cuda_device_with_base_gpu_id(self): mock_engine_class.from_pretrained.return_value = MagicMock() with patch( - "torchspec.inference.engine.hf_runner.HFRunner", + "aurora.inference.engine.hf_runner.HFRunner", mock_engine_class, ): - with patch("torchspec.config.mooncake_config.MooncakeConfig", MagicMock()): + with patch("aurora.config.mooncake_config.MooncakeConfig", MagicMock()): with patch("torch.cuda.set_device") as mock_set_device: - with patch("torchspec.ray.ray_actor._to_local_gpu_id", return_value=0): + with patch("aurora.ray.ray_actor._to_local_gpu_id", return_value=0): engine.init(mooncake_config=None) mock_set_device.assert_called_once_with(0) diff --git a/tests/test_hf_runner.py b/tests/test_hf_runner.py index 2eb5fd7..5da45cd 100644 --- a/tests/test_hf_runner.py +++ b/tests/test_hf_runner.py @@ -4,9 +4,9 @@ import torch -from torchspec.config.inference_config import HFInferenceConfig -from torchspec.config.mooncake_config import MooncakeConfig -from torchspec.inference.engine.hf_runner import HFRunner +from aurora.config.inference_config import HFInferenceConfig +from aurora.config.mooncake_config import MooncakeConfig +from aurora.inference.engine.hf_runner import HFRunner class MockMooncakeStore: @@ -54,7 +54,7 @@ def test_init_mooncake_store_with_config(self): assert engine.mooncake_store is None with patch( - "torchspec.inference.engine.hf_runner.EagleMooncakeStore", + "aurora.inference.engine.hf_runner.EagleMooncakeStore", MockMooncakeStore, ): with patch("torch.cuda.current_device", return_value=0): @@ -74,7 +74,7 @@ def test_init_mooncake_store_with_explicit_config(self): ) with patch( - "torchspec.inference.engine.hf_runner.EagleMooncakeStore", + "aurora.inference.engine.hf_runner.EagleMooncakeStore", MockMooncakeStore, ): with patch("torch.cuda.current_device", return_value=0): @@ -108,7 +108,7 @@ def test_setup_initializes_mooncake_store_if_configured(self): with patch.object(engine, "_setup_target_model"): with patch( - "torchspec.inference.engine.hf_runner.EagleMooncakeStore", + "aurora.inference.engine.hf_runner.EagleMooncakeStore", MockMooncakeStore, ): with patch("torch.cuda.current_device", return_value=0): diff --git a/tests/test_mooncake_master.py b/tests/test_mooncake_master.py index b070a2b..76c5a93 100644 --- a/tests/test_mooncake_master.py +++ b/tests/test_mooncake_master.py @@ -6,7 +6,7 @@ from argparse import Namespace from unittest import mock -from torchspec.transfer.mooncake.utils import ( +from aurora.transfer.mooncake.utils import ( MooncakeMaster, launch_mooncake_master, resolve_mooncake_master_bin, @@ -45,13 +45,13 @@ def test_start_launches_subprocess(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/usr/bin/mooncake_master", ), mock.patch("os.path.exists", return_value=True), mock.patch("subprocess.Popen", return_value=mock_process) as mock_popen, mock.patch("time.sleep"), - mock.patch("torchspec.ray.ray_actor.get_current_node_ip", return_value="10.0.0.1"), + mock.patch("aurora.ray.ray_actor.get_current_node_ip", return_value="10.0.0.1"), ): info = actor.start(50051, 8090, "0.0.0.0") @@ -68,7 +68,7 @@ def test_start_raises_on_missing_binary(self): actor = MooncakeMaster() with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/nonexistent/mooncake_master", ), mock.patch("os.path.exists", return_value=False), @@ -89,7 +89,7 @@ def test_start_raises_on_process_failure(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/usr/bin/mooncake_master", ), mock.patch("os.path.exists", return_value=True), @@ -178,7 +178,7 @@ def test_creates_named_actor_and_writes_back_to_args(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/usr/bin/mooncake_master", ), mock.patch("os.path.exists", return_value=True), @@ -210,13 +210,13 @@ def test_auto_resolves_addr_and_port(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/usr/bin/mooncake_master", ), mock.patch("os.path.exists", return_value=True), - mock.patch("torchspec.ray.ray_actor.get_current_node_ip", return_value="10.0.0.1"), + mock.patch("aurora.ray.ray_actor.get_current_node_ip", return_value="10.0.0.1"), mock.patch( - "torchspec.ray.ray_actor.get_free_port", + "aurora.ray.ray_actor.get_free_port", side_effect=[55000, 8500], ), mock.patch("ray.remote", return_value=mock_decorator), @@ -237,7 +237,7 @@ def test_returns_none_when_binary_not_found(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/nonexistent/path/mooncake_master", ), mock.patch("os.path.exists", return_value=False), @@ -258,7 +258,7 @@ def test_returns_none_on_start_failure(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value="/usr/bin/mooncake_master", ), mock.patch("os.path.exists", return_value=True), @@ -288,10 +288,10 @@ def test_launches_real_subprocess(self): with ( mock.patch( - "torchspec.transfer.mooncake.utils.resolve_mooncake_master_bin", + "aurora.transfer.mooncake.utils.resolve_mooncake_master_bin", return_value=script_path, ), - mock.patch("torchspec.ray.ray_actor.get_current_node_ip", return_value="127.0.0.1"), + mock.patch("aurora.ray.ray_actor.get_current_node_ip", return_value="127.0.0.1"), ): info = actor.start(50051, 8090, "0.0.0.0") diff --git a/tests/test_vocab_mapping.py b/tests/test_vocab_mapping.py index 8d742da..a2237b9 100644 --- a/tests/test_vocab_mapping.py +++ b/tests/test_vocab_mapping.py @@ -1,4 +1,4 @@ -"""Test generate_vocab_mapping and _count_token_frequencies from torchspec.data.preprocessing. +"""Test generate_vocab_mapping and _count_token_frequencies from aurora.data.preprocessing. Verifies that: - Token counting via packed_loss_mask matches a reference implementation. @@ -11,12 +11,12 @@ import pytest import torch -from torchspec.data.preprocessing import ( +from aurora.data.preprocessing import ( _count_token_frequencies, generate_vocab_mapping, process_token_dict_to_mappings, ) -from torchspec.data.utils import pack_loss_mask, serialize_packed_loss_mask +from aurora.data.utils import pack_loss_mask, serialize_packed_loss_mask def _make_prompts(input_ids_list, loss_mask_list): diff --git a/tools/bench_eagle3_mask_modes.py b/tools/bench_eagle3_mask_modes.py index af624d2..919f7f8 100755 --- a/tools/bench_eagle3_mask_modes.py +++ b/tools/bench_eagle3_mask_modes.py @@ -42,7 +42,7 @@ import torch -from torchspec.models.draft.llama3_eagle import ( +from aurora.models.draft.llama3_eagle import ( _build_eagle3_mask_pair, _EagleMaskedFlashAttnFunc, _snap_q_len, @@ -58,7 +58,7 @@ def _clear_compile_caches(): """Clear all flash_attn compile caches + mask_mod caches.""" - from torchspec.models.draft import llama3_eagle as mod + from aurora.models.draft import llama3_eagle as mod # mask_mod caches mod._flash_mask_mod_cache.clear() diff --git a/tools/benchmark_eagle3.py b/tools/benchmark_eagle3.py index 6ed61f3..e43b3fb 100755 --- a/tools/benchmark_eagle3.py +++ b/tools/benchmark_eagle3.py @@ -19,13 +19,13 @@ import torch from transformers.models.llama.configuration_llama import LlamaConfig -from torchspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3 -from torchspec.models.eagle3 import ( +from aurora.models.draft.llama3_eagle import LlamaForCausalLMEagle3 +from aurora.models.eagle3 import ( Eagle3Model, compute_lazy_target_padded, compute_target_p_padded, ) -from torchspec.training.optimizer import BF16Optimizer +from aurora.training.optimizer import BF16Optimizer # --------------------------------------------------------------------------- diff --git a/tools/build_conda.sh b/tools/build_conda.sh index ae160a9..cbfd92d 100755 --- a/tools/build_conda.sh +++ b/tools/build_conda.sh @@ -58,7 +58,7 @@ elif [ "$MODE" = "current" ]; then pip install -e "${SGLANG_FOLDER_NAME}/python[all]" pip install -e ".[dev]" - echo "torchspec installed into current environment!" + echo "aurora installed into current environment!" else echo "Skipping package installation (mode=0)" echo "Please install packages manually:" diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py index 1b92297..4a9dcbc 100644 --- a/tools/convert_to_hf.py +++ b/tools/convert_to_hf.py @@ -56,7 +56,7 @@ from tqdm import tqdm from typing_extensions import override -from torchspec.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel +from aurora.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel logging.basicConfig( level=logging.INFO, @@ -124,7 +124,7 @@ def _detect_model_dir(input_dir: str) -> str: def _generate_config_from_target( target_model_path: str, output_path: str, trust_remote_code: bool = False ) -> str: - from torchspec.config.utils import generate_draft_model_config + from aurora.config.utils import generate_draft_model_config logger.info("Auto-generating draft model config from %s", target_model_path) config_dict = generate_draft_model_config( @@ -263,7 +263,7 @@ def _count_token_frequencies(prompts: list[dict]) -> Counter: import numba import numpy as np - from torchspec.data.utils import unpack_loss_mask + from aurora.data.utils import unpack_loss_mask @numba.njit(cache=True) def _histogram(ids, mask, counts): @@ -304,7 +304,7 @@ def _load_tokenized_prompts( max_seq_length: int, cache_dir: Optional[str], ) -> list: - from torchspec.data.dataset import load_conversation_dataset + from aurora.data.dataset import load_conversation_dataset args_ns = argparse.Namespace( train_data_path=dataset_path, @@ -408,7 +408,7 @@ def _convert_fsdp_to_hf( # ── Vocab pruning ──────────────────────────────────────────────────── assert dataset_path is not None and draft_vocab_size is not None assert tokenizer is not None and chat_template is not None - from torchspec.data.preprocessing import process_token_dict_to_mappings + from aurora.data.preprocessing import process_token_dict_to_mappings logger.info( "Vocab pruning: vocab_size=%d, draft_vocab_size=%d", @@ -552,7 +552,7 @@ def _validate_args(args: argparse.Namespace) -> None: raise ValueError("--draft-vocab-size is required when --prune-vocab is set") if args.chat_template: - from torchspec.data.template import TEMPLATE_REGISTRY + from aurora.data.template import TEMPLATE_REGISTRY available = TEMPLATE_REGISTRY.get_all_template_names() if args.chat_template not in available: diff --git a/tools/kill_all_torchspec.sh b/tools/kill_all_aurora.sh similarity index 75% rename from tools/kill_all_torchspec.sh rename to tools/kill_all_aurora.sh index 8574b25..7ff4528 100755 --- a/tools/kill_all_torchspec.sh +++ b/tools/kill_all_aurora.sh @@ -3,16 +3,16 @@ if [ "$1" = "rocm" ]; then echo "Running in ROCm mode" - echo "Killing torchspec inference workers..." - pgrep -f 'torchspec\.target\.remote_backend' | xargs -r kill -9 + echo "Killing aurora inference workers..." + pgrep -f 'aurora\.target\.remote_backend' | xargs -r kill -9 echo "Killing mooncake master..." pgrep -f 'mooncake_master' | xargs -r kill -9 else nvidia-smi - echo "Killing torchspec inference workers..." - pgrep -f 'torchspec\.target\.remote_backend' | xargs -r kill -9 + echo "Killing aurora inference workers..." + pgrep -f 'aurora\.target\.remote_backend' | xargs -r kill -9 echo "Killing mooncake master..." pgrep -f 'mooncake_master' | xargs -r kill -9 diff --git a/tools/max_seq_search.py b/tools/max_seq_search.py index 50ada6c..29e9a98 100755 --- a/tools/max_seq_search.py +++ b/tools/max_seq_search.py @@ -24,12 +24,12 @@ import torch import torch.nn as nn -from torchspec import AutoDraftModelConfig, AutoEagle3DraftModel, Eagle3Model -from torchspec.config.utils import generate_draft_model_config -from torchspec.models.eagle3 import compute_lazy_target_padded, compute_target_p_padded -from torchspec.models.target.target_utils import TargetLMHead -from torchspec.training.optimizer import BF16Optimizer -from torchspec.utils.memory import available_memory, clear_memory +from aurora import AutoDraftModelConfig, AutoEagle3DraftModel, Eagle3Model +from aurora.config.utils import generate_draft_model_config +from aurora.models.eagle3 import compute_lazy_target_padded, compute_target_p_padded +from aurora.models.target.target_utils import TargetLMHead +from aurora.training.optimizer import BF16Optimizer +from aurora.utils.memory import available_memory, clear_memory def create_synthetic_batch( diff --git a/tools/test_sglang_engine_patch.py b/tools/test_sglang_engine_patch.py index 94d1b5e..5463032 100755 --- a/tools/test_sglang_engine_patch.py +++ b/tools/test_sglang_engine_patch.py @@ -122,8 +122,8 @@ def set_mooncake_env(host, grpc_port, http_port): def test_spec_training(model_path, aux_layer_ids=None): import sglang as sgl - from torchspec.config.mooncake_config import MooncakeConfig - from torchspec.transfer.mooncake import EagleMooncakeStore + from aurora.config.mooncake_config import MooncakeConfig + from aurora.transfer.mooncake import EagleMooncakeStore if aux_layer_ids is None: from transformers import AutoConfig diff --git a/tools/update_sglang_patch.sh b/tools/update_sglang_patch.sh index 4d2f2f1..95f91c2 100755 --- a/tools/update_sglang_patch.sh +++ b/tools/update_sglang_patch.sh @@ -95,7 +95,7 @@ echo "Generating patch from $SGLANG_COMMIT to HEAD..." # git apply ignores lines before the first "diff --git" line, # so the diffstat is purely informational for human readers. { - echo "torchspec sglang patch (base: ${SGLANG_COMMIT:0:10})" + echo "aurora sglang patch (base: ${SGLANG_COMMIT:0:10})" echo "---" git diff --stat "$SGLANG_COMMIT" HEAD echo ""