tuft 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

tuft/backends/hf_training_model.py

@@ -1,6 +1,7 @@
 import asyncio
+import logging
 import shutil
-from typing import Dict
+from typing import Callable, Dict
 
 import ray
 import torch
@@ -14,13 +15,12 @@ from transformers import AutoModelForCausalLM
 
 from tuft.checkpoints import CheckpointRecord
 from tuft.config import ModelConfig
-from tuft.loss_fn import get_loss_fn
+from tuft.loss_fn import get_loss_fn, metrics_reduction
 from tuft.telemetry.tracing import extract_context, get_tracer
 
 
 _get_tracer = lambda: get_tracer("tuft.hf_training_model")  # noqa: E731
 
-
 MODULE_MAP = {
     "llama": {
         "attn": ["q_proj", "k_proj", "v_proj", "o_proj"],
@@ -58,6 +58,8 @@ class HFTrainingModel:
         self.model = self._init_peft_model(config)
         self.adapter_optimizer: Dict[str, torch.optim.AdamW] = {}
         self._lock = asyncio.Lock()
+        self.logger = logging.getLogger()
+        self.micro_batch_size = config.micro_batch_size
 
     async def async_init(self) -> None:
         """Do nothing for now. Just used to make sure the actor is ready."""
@@ -193,7 +195,9 @@ class HFTrainingModel:
         async with self._lock:
             if lora_id in self.adapter_optimizer:
                 self.model.delete_adapter(lora_id)
-                self.adapter_optimizer.pop(lora_id)
+                optimizer = self.adapter_optimizer.pop(lora_id)
+                del optimizer
+                torch.cuda.empty_cache()
 
     # --------------------------------
     # Training methods
@@ -207,7 +211,7 @@ class HFTrainingModel:
         backward: bool = False,
         trace_context: dict[str, str] | None = None,
     ) -> types.ForwardBackwardOutput:
-        """Forward pass (and backward if specified).
+        """Forward pass with micro-batch gradient accumulation.
 
         Args:
             data: List of Datum objects containing input data.
@@ -222,73 +226,163 @@ class HFTrainingModel:
         """
         ctx = extract_context(trace_context or {})
         span_name = "hf_model.forward_backward" if backward else "hf_model.forward"
+
         with _get_tracer().start_as_current_span(span_name, context=ctx) as span:
             span.set_attribute("tuft.lora_id", lora_id)
             span.set_attribute("tuft.backward", backward)
             span.set_attribute("tuft.data_count", len(data))
-            # Prepare input tensors
-            input_ids = [
-                torch.tensor(datum.model_input.to_ints(), dtype=torch.long) for datum in data
-            ]
-            input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=0)
-            attention_mask = (input_ids_padded != 0).long()
-            position_ids = (
-                torch.arange(input_ids_padded.size(1), dtype=torch.long)
-                .unsqueeze(0)
-                .expand(input_ids_padded.size(0), -1)
-            )
-            # Move tensors to model device
-            device = next(self.model.parameters()).device
-            input_ids_padded = input_ids_padded.to(device)
-            attention_mask = attention_mask.to(device)
-            position_ids = position_ids.to(device)
-
-            # Activate the correct adapter
-            async with self._lock:
-                self._activate_adapter(lora_id)
-
-                # Forward pass
-                outputs = self.model(
-                    input_ids=input_ids_padded,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    return_dict=True,
-                )
 
-                # Compute loss
-                if loss_fn_config is None:
-                    loss_fn_config = {}
-                loss_fn_callable = get_loss_fn(loss_fn)
-                logits = outputs.logits
-                if "temperature" in loss_fn_config:
-                    temperature = loss_fn_config["temperature"]
-                    logits.div_(temperature)
+            batch_size = len(data)
+            micro_batch_size = self.config.micro_batch_size
 
-                loss_fn_inputs = self._prepare_loss_fn_inputs(data)
+            num_micro_batches = (batch_size + micro_batch_size - 1) // micro_batch_size
+            span.set_attribute("tuft.num_micro_batches", num_micro_batches)
 
-                ## compute target_logprobs from logits and target_tokens
-                target_tokens = loss_fn_inputs["target_tokens"]
-                target_logprobs = self._compute_logprobs_from_target_tokens(logits, target_tokens)
-                loss_fn_inputs["target_logprobs"] = target_logprobs
+            if num_micro_batches > 1:
+                self.logger.debug(
+                    f"[MICRO_BATCH] Splitting batch_size={batch_size} into "
+                    f"{num_micro_batches} micro-batches of size {micro_batch_size}"
+                )
 
-                loss, metric = loss_fn_callable(loss_fn_inputs, loss_fn_config)
+            loss_fn_callable = get_loss_fn(loss_fn)
+            all_loss_fn_outputs = []
+            micro_batch_weights = []
+            metric_list = []
+            total_loss = 0.0
 
-                # Backward pass if needed
-                if backward:
-                    loss.backward()
+            async with self._lock:
+                self._activate_adapter(lora_id)
 
-            unpaded_logprobs = self._unpad_tensor(
-                target_logprobs, [len(datum.model_input.to_ints()) for datum in data]
+                for micro_idx in range(num_micro_batches):
+                    start_idx = micro_idx * micro_batch_size
+                    end_idx = min(start_idx + micro_batch_size, batch_size)
+                    micro_data = data[start_idx:end_idx]
+
+                    torch.cuda.reset_peak_memory_stats()
+                    self.logger.debug(
+                        f"[GPU-micro_batch_{micro_idx}] before_forward: "
+                        f"allocated={torch.cuda.memory_allocated() / 1e9:.2f}GB, "
+                        f"reserved={torch.cuda.memory_reserved() / 1e9:.2f}GB"
+                    )
+
+                    micro_loss, micro_metrics, micro_outputs = await self._forward_micro_batch(
+                        micro_data,
+                        loss_fn_callable,
+                        loss_fn_config,
+                        backward=backward,
+                    )
+
+                    total_loss += micro_loss
+                    all_loss_fn_outputs.extend(micro_outputs)
+                    micro_batch_weights.append(len(micro_outputs))
+
+                    metric_list.append(micro_metrics)
+
+                    self.logger.debug(
+                        f"[GPU-micro_batch_{micro_idx}] after_forward: "
+                        f"allocated={torch.cuda.memory_allocated() / 1e9:.2f}GB, "
+                        f"reserved={torch.cuda.memory_reserved() / 1e9:.2f}GB, "
+                        f"max_allocated={torch.cuda.max_memory_allocated() / 1e9:.2f}GB"
+                    )
+
+                    torch.cuda.empty_cache()
+
+            avg_loss = total_loss / num_micro_batches
+            self.logger.debug(f"Average loss: {avg_loss}")
+            metric_list = metrics_reduction(metric_list, micro_batch_weights)
+
+            self.logger.debug(
+                f"[GPU-after_micro_batches] allocated={torch.cuda.memory_allocated() / 1e9:.2f}GB"
+                f", reserved={torch.cuda.memory_reserved() / 1e9:.2f}GB"
             )
+
             return types.ForwardBackwardOutput(
                 loss_fn_output_type=loss_fn,
-                loss_fn_outputs=[
-                    {"logprobs": types.TensorData.from_torch(logprobs.detach())}
-                    for logprobs in unpaded_logprobs
-                ],
-                metrics=metric,
+                loss_fn_outputs=all_loss_fn_outputs,
+                metrics=metric_list or {},
             )
 
+    async def _forward_micro_batch(
+        self,
+        data: list[types.Datum],
+        loss_fn_callable: Callable,
+        loss_fn_config: dict[str, float] | None,
+        backward: bool,
+    ) -> tuple[float, dict[str, float], list[dict]]:
+        """Process a single micro-batch.
+
+        Returns:
+            tuple: (loss_value, metrics_dict, loss_fn_outputs_list)
+        """
+        # Prepare input tensors
+        input_ids = [torch.tensor(datum.model_input.to_ints(), dtype=torch.long) for datum in data]
+        input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=0)
+        attention_mask = (input_ids_padded != 0).long()
+        position_ids = (
+            torch.arange(input_ids_padded.size(1), dtype=torch.long)
+            .unsqueeze(0)
+            .expand(input_ids_padded.size(0), -1)
+        )
+
+        device = next(self.model.parameters()).device
+        input_ids_padded = input_ids_padded.to(device)
+        attention_mask = attention_mask.to(device)
+        position_ids = position_ids.to(device)
+
+        # Forward pass
+        outputs = self.model(
+            input_ids=input_ids_padded,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            return_dict=True,
+        )
+
+        if loss_fn_config is None:
+            loss_fn_config = {}
+
+        logits = outputs.logits
+        del outputs
+        torch.cuda.empty_cache()
+
+        if "temperature" in loss_fn_config:
+            temperature = loss_fn_config["temperature"]
+            logits = logits / temperature
+
+        loss_fn_inputs = self._prepare_loss_fn_inputs(data)
+        target_tokens = loss_fn_inputs["target_tokens"]
+
+        target_logprobs = self._compute_logprobs_from_target_tokens(logits, target_tokens)
+        del logits
+        torch.cuda.empty_cache()
+
+        loss_fn_inputs["target_logprobs"] = target_logprobs
+        loss, metric = loss_fn_callable(loss_fn_inputs, loss_fn_config)
+
+        # Backward with gradient accumulation
+        if backward:
+            loss.backward(retain_graph=False)
+            torch.cuda.empty_cache()
+
+        unpaded_logprobs = self._unpad_tensor(
+            target_logprobs.detach(),
+            [len(datum.model_input.to_ints()) for datum in data],
+        )
+        loss_fn_outputs = [
+            {"logprobs": types.TensorData.from_torch(logprobs.cpu().clone())}
+            for logprobs in unpaded_logprobs
+        ]
+
+        loss_value = loss.detach().item()
+
+        del target_logprobs
+        del unpaded_logprobs
+        del loss_fn_inputs
+        del loss
+
+        torch.cuda.empty_cache()
+
+        return loss_value, metric, loss_fn_outputs
+
     async def optim_step(
         self,
         adam_params: types.AdamParams,
@@ -316,7 +410,9 @@ class HFTrainingModel:
                 param_group["weight_decay"] = adam_params.weight_decay
             optimizer.step()
             optimizer.zero_grad()
-            return types.OptimStepResponse()
+
+        torch.cuda.empty_cache()
+        return types.OptimStepResponse()
 
     # --------------------------------
     # Helper methods
@@ -362,11 +458,33 @@ class HFTrainingModel:
     def _compute_logprobs_from_target_tokens(
         self, logits: torch.Tensor, target_tokens: torch.Tensor
     ) -> torch.Tensor:
-        logits_labels = torch.gather(logits, dim=-1, index=target_tokens.unsqueeze(-1)).squeeze(-1)
-        # loop to reduce peak mem consumption
-        logsumexp_values = torch.stack([torch.logsumexp(logit, dim=-1) for logit in logits])
-        logprobs_labels = logits_labels - logsumexp_values  # log_softmax(x_i) = x_i - logsumexp(x)
-        return logprobs_labels
+        """Compute log probabilities of target tokens from logits with low memory usage.
+        https://github.com/OpenRLHF/OpenRLHF/pull/718
+        """
+        if logits.dtype in [torch.float32, torch.float64]:
+            logits_labels = torch.gather(logits, dim=-1, index=target_tokens.unsqueeze(-1)).squeeze(
+                -1
+            )
+            logsumexp_values = torch.stack(
+                [
+                    torch.logsumexp(logit, dim=-1) for logit in logits
+                ]  # loop to reduce peak mem consumption
+            )
+            log_probs_labels = (
+                logits_labels - logsumexp_values
+            )  # log_softmax(x_i) = x_i - logsumexp(x)
+        else:
+            log_probs_labels = []
+            for row_logits, row_labels in zip(
+                logits, target_tokens, strict=True
+            ):  # loop to reduce peak mem consumption
+                row_log_probs = torch.nn.functional.log_softmax(row_logits, dim=-1)
+                row_log_probs_labels = row_log_probs.gather(
+                    dim=-1, index=row_labels.unsqueeze(-1)
+                ).squeeze(-1)
+                log_probs_labels.append(row_log_probs_labels)
+            log_probs_labels = torch.stack(log_probs_labels)
+        return log_probs_labels
 
     def _unpad_tensor(
         self, padded_tensor: torch.Tensor, original_lengths: list[int]
@@ -383,6 +501,8 @@ class HFTrainingModel:
             dtype="auto",
             device_map="auto",
         )
+        model.enable_input_require_grads()
+        model.gradient_checkpointing_enable({"use_reentrant": False})
         peft_config = LoraConfig()
         peft_model = get_peft_model(model, peft_config=peft_config, adapter_name="default")
         return peft_model
@@ -398,7 +518,7 @@ class HFTrainingModel:
             ray.remote(cls)
             .options(
                 name="training_model_" + config.model_name,
-                num_gpus=1 if not config.colocate else 1 - config.sampling_memory_fraction,
+                num_gpus=(1 if not config.colocate else 1 - config.sampling_memory_fraction),
             )
             .remote(config)
         )

tuft/cli.py

@@ -10,12 +10,21 @@ import typer
 import uvicorn
 
 from .config import AppConfig, load_yaml_config
+from .exceptions import ConfigMismatchError
+from .persistence import (
+    flush_all_data,
+    get_current_namespace,
+    get_redis_store,
+    validate_config_signature,
+)
 from .server import create_root_app
 from .telemetry import init_telemetry
 from .telemetry.metrics import ResourceMetricsCollector
 
 
 app = typer.Typer(help="TuFT - Tenant-unified Fine-Tuning Server.", no_args_is_help=True)
+clear_app = typer.Typer(help="Clear data commands.", no_args_is_help=True)
+app.add_typer(clear_app, name="clear")
 
 
 # Required for Typer to recognize subcommands when using no_args_is_help=True
@@ -79,6 +88,113 @@ def _build_config(
     return config
 
 
+_FORCE_OPTION = typer.Option(
+    False,
+    "--force",
+    "-f",
+    help="Skip confirmation prompts when clearing persistence data.",
+)
+
+
+@clear_app.command(name="persistence")
+def clear_persistence(
+    config_path: Path | None = _CONFIG_OPTION,
+    force: bool = _FORCE_OPTION,
+) -> None:
+    """Clear all persistence data and start fresh.
+
+    This command clears all existing persistence data in the configured namespace.
+    Use this when the configuration has changed and you want to discard old data.
+    """
+    # Build config to get persistence settings
+    try:
+        resolved_config_path = _resolve_config_path(config_path)
+        config = load_yaml_config(resolved_config_path)
+    except typer.BadParameter as e:
+        typer.secho(f"Error: {e}", fg=typer.colors.RED)
+        raise typer.Exit(1) from e
+
+    if not config.persistence.enabled:
+        typer.secho(
+            "Persistence is disabled in the configuration. Nothing to clear.",
+            fg=typer.colors.YELLOW,
+        )
+        raise typer.Exit(0)
+
+    # Configure the store
+    store = get_redis_store()
+    store.configure(config.persistence)
+    namespace = get_current_namespace()
+
+    if not force:
+        typer.secho(
+            "\n🚨🚨🚨 CRITICAL WARNING 🚨🚨🚨\n",
+            fg=typer.colors.RED,
+            bold=True,
+        )
+        typer.secho(
+            "This command will PERMANENTLY DELETE ALL persistence data!\n",
+            fg=typer.colors.RED,
+            bold=True,
+        )
+        typer.secho(
+            f"📦 Target namespace: '{namespace}'\n",
+            fg=typer.colors.YELLOW,
+            bold=True,
+        )
+        typer.echo(
+            f"This IRREVERSIBLE action will destroy ALL data in namespace '{namespace}':\n"
+            "  ❌ All saved sessions\n"
+            "  ❌ All training run records and checkpoint metadata (NOT local checkpoint files)\n"
+            "  ❌ All future records\n"
+            "  ❌ All sampling session records\n"
+            "  ❌ Configuration signature\n"
+            "\n"
+            "⚠️  The server will start fresh with NO previous state.\n"
+            "⚠️  This action CANNOT be undone!\n"
+            "⚠️  Local checkpoint files on disk are NOT affected.\n"
+            f"⚠️  Only data in namespace '{namespace}' will be affected.\n"
+        )
+        confirmed = typer.confirm(
+            f"Do you REALLY want to delete all data in namespace '{namespace}'?",
+            default=False,
+        )
+        if not confirmed:
+            typer.echo("Aborted. No data was cleared.")
+            raise typer.Exit(0)
+
+    deleted_count, cleared_namespace = flush_all_data()
+    typer.secho(
+        f"✅ Cleared {deleted_count} keys from namespace '{cleared_namespace}'.",
+        fg=typer.colors.GREEN,
+    )
+    typer.echo("Persistence data has been cleared. You can now start the server fresh.")
+
+
+def _validate_persistence_config(config: AppConfig) -> None:
+    """Validate that persistence config matches stored config.
+
+    If config mismatch is detected, exits with an error message.
+    """
+    if not config.persistence.enabled:
+        return
+
+    # Configure the Redis store first
+    store = get_redis_store()
+    store.configure(config.persistence)
+
+    try:
+        validate_config_signature(config)
+    except ConfigMismatchError as e:
+        typer.secho(
+            "\n 🚫 FATAL ERROR: Configuration Mismatch Detected 🚫",
+            fg=typer.colors.RED,
+            bold=True,
+        )
+        typer.echo(f"\n{e}\n")
+        raise typer.Exit(1) from e
+
+
 def _init_telemetry(config: AppConfig, log_level: str) -> None:
     """Initialize OpenTelemetry if enabled."""
     # Configure root logger level to ensure logs flow to OTel
@@ -104,6 +220,10 @@ def launch(
 ) -> None:
     """Launch the TuFT server."""
     app_config = _build_config(config_path, checkpoint_dir)
+
+    # Validate persistence configuration before starting
+    _validate_persistence_config(app_config)
+
     # Initialize telemetry before starting the server
     _init_telemetry(app_config, log_level)
     logging.getLogger("tuft").info("Server starting on %s:%s", host, port)

tuft/config.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, Iterable, List
+from typing import Any, Iterable
+
+from pydantic import BaseModel, Field, model_validator
 
 from .persistence import PersistenceConfig
 
@@ -14,12 +15,7 @@ def _default_checkpoint_dir() -> Path | None:
     return None
 
 
-def _default_persistence_config() -> PersistenceConfig:
-    return PersistenceConfig()
-
-
-@dataclass
-class TelemetryConfig:
+class TelemetryConfig(BaseModel):
     """Configuration for OpenTelemetry integration.
 
     Attributes:
@@ -32,26 +28,61 @@ class TelemetryConfig:
     enabled: bool = False
     service_name: str = "tuft"
     otlp_endpoint: str | None = None
-    resource_attributes: Dict[str, str] = field(default_factory=dict)
+    resource_attributes: dict[str, str] = Field(default_factory=dict)
+
+
+class ModelConfig(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+
+    model_name: str  # name used in APIs
+    model_path: Path  # path to model checkpoint
+    max_model_len: int  # maximum context length supported by the model
+    tensor_parallel_size: int = 1  # tensor parallel size
+
+    # default sampling parameters for this model
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = -1
+    logprobs: int = 0
+    seed: int = 42
+    min_response_tokens: int = 0
+
+    # default lora setting
+    max_lora_rank: int = 16  # maximum rank for LoRA adapters
+    max_loras: int = 1  # maximum number of LoRA adapters that can be applied simultaneously
 
+    # default training setting
+    micro_batch_size: int = 1  # micro-batch size for training
 
-def _default_telemetry_config() -> TelemetryConfig:
-    return TelemetryConfig()
+    # whether to colocate sampling and training on the same device
+    # only for local testing purposes
+    colocate: bool = False
+    sampling_memory_fraction: float = 0.2  # fraction of GPU memory for sampling
+
+    @model_validator(mode="after")
+    def validate_colocate(self) -> "ModelConfig":
+        if self.colocate and self.tensor_parallel_size != 1:
+            raise ValueError("Colocate option is only supported for tensor_parallel_size=1.")
+        return self
 
 
-@dataclass
-class AppConfig:
-    """Runtime configuration for the TuFT server."""
+class AppConfig(BaseModel):
+    """Runtime configuration for the TuFT server.
 
-    checkpoint_dir: Path | None = field(default_factory=_default_checkpoint_dir)
-    supported_models: List[ModelConfig] = field(default_factory=list)
+    This is a Pydantic model that can be serialized/deserialized for persistence.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    checkpoint_dir: Path | None = Field(default_factory=_default_checkpoint_dir)
+    supported_models: list[ModelConfig] = Field(default_factory=list)
     model_owner: str = "local-user"
     toy_backend_seed: int = 0
     # TODO: Temporary implementation for user authorization,
     # replace with proper auth system later
-    authorized_users: Dict[str, str] = field(default_factory=dict)
-    persistence: PersistenceConfig = field(default_factory=_default_persistence_config)
-    telemetry: TelemetryConfig = field(default_factory=_default_telemetry_config)
+    authorized_users: dict[str, str] = Field(default_factory=dict)
+    persistence: PersistenceConfig = Field(default_factory=PersistenceConfig)
+    telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig)
 
     def ensure_directories(self) -> None:
         if self.checkpoint_dir is not None:
@@ -74,50 +105,21 @@ class AppConfig:
             self.supported_models = updated
         return self
 
-
-@dataclass
-class ModelConfig:
-    """Configuration for a specific model."""
-
-    model_name: str  # name used in APIs
-    model_path: Path  # path to model checkpoint
-    max_model_len: int  # maximum context length supported by the model
-    tensor_parallel_size: int = 1  # tensor parallel size
-
-    # default sampling parameters for this model
-    temperature: float = 1.0
-    top_p: float = 1.0
-    top_k: int = -1
-    logprobs: int = 0
-    seed: int = 42
-    min_response_tokens: int = 0
-
-    # default lora setting
-    max_lora_rank: int = 16  # maximum rank for LoRA adapters
-    max_loras: int = 1  # maximum number of LoRA adapters that can be applied simultaneously
-
-    # whether to colocate sampling and training on the same device
-    # only for local testing purposes
-    colocate: bool = False
-    sampling_memory_fraction: float = 0.2  # fraction of GPU memory for sampling
-
-    def __post_init__(self) -> None:
-        if self.colocate and self.tensor_parallel_size != 1:
-            raise ValueError("Colocate option is only supported for tensor_parallel_size=1.")
+    def get_config_for_persistence(self) -> dict[str, Any]:
+        """Get config fields for persistence signature (excludes persistence config itself)."""
+        return self.model_dump(mode="json", exclude={"persistence"})
 
 
 def load_yaml_config(config_path: Path) -> AppConfig:
     """Loads an AppConfig from a YAML file."""
     from omegaconf import OmegaConf
 
-    schema = OmegaConf.structured(AppConfig)
     loaded = OmegaConf.load(config_path)
     try:
-        config = OmegaConf.merge(schema, loaded)
-        app_config = OmegaConf.to_object(config)
-        assert isinstance(app_config, AppConfig), (
-            "Loaded config is not of type AppConfig, which should not happen."
-        )
-        return app_config
+        # Convert OmegaConf to plain dict for Pydantic
+        config_dict = OmegaConf.to_container(loaded, resolve=True)
+        if not isinstance(config_dict, dict):
+            raise ValueError("Config file must contain a dictionary at root level")
+        return AppConfig.model_validate(config_dict)
     except Exception as e:
         raise ValueError(f"Failed to load config from {config_path}: {e}") from e