d9d 0.1.0__py3-none-any.whl

d9d/loop/component/job_logger.py

@@ -0,0 +1,146 @@
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+import torch
+import torch.utils._pytree as pytree  # noqa: PLC2701
+from torch.distributed.checkpoint.stateful import Stateful
+
+from d9d.core.dist_context import DistributedContext
+from d9d.core.types import PyTree, ScalarTree
+from d9d.internals.state import load_state_dict_main_process, state_dict_main_process
+from d9d.loop.config import JobLoggerConfig
+from d9d.metric.impl import ComposeMetric
+from d9d.tracker import BaseTracker, BaseTrackerRun, RunConfig, tracker_from_config
+from d9d.tracker.provider.null import NullTrackerConfig
+
+from .stepper import Stepper
+
+
+def _flatten_pytree_for_metrics(tree: PyTree[float]) -> dict[str, float]:
+    flat_dict = {}
+
+    for path_tuple, value in pytree.tree_leaves_with_path(tree):
+        path_segments = []
+
+        for key in path_tuple:
+            match key:
+                case pytree.MappingKey(k):
+                    path_segments.append(str(k))
+                case pytree.SequenceKey(idx):
+                    path_segments.append(str(idx))
+                case pytree.GetAttrKey(name):
+                    path_segments.append(name)
+                case _:
+                    path_segments.append(str(key))
+
+        flat_key = "/".join(path_segments)
+        flat_dict[flat_key] = value
+
+    return flat_dict
+
+
+class JobLogger(Stateful):
+    """
+    Handles the logging of training metrics and loss values.
+
+    This class coordinates with the distributed context and metric calculators
+    to log instantaneous loss values and periodic aggregated metrics to the
+    configured experiment tracker.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            config: JobLoggerConfig,
+            metrics: ComposeMetric,
+            stepper: Stepper,
+            run_config: RunConfig,
+            additional_hparams: ScalarTree
+    ):
+        """
+        Constructs JobLogger object.
+
+        Args:
+            dist_context: The distributed context.
+            config: Configuration settings.
+            metrics: The composite metric collection to be computed and logged.
+            stepper: Object tracking the current global step.
+            run_config: Run configuration.
+        """
+
+        self._dist_context = dist_context
+        self._config = config
+        self._metrics = metrics
+        self._stepper = stepper
+        self._run_config = run_config.model_copy(deep=True, update={"hparams": {
+            "run": run_config.hparams,
+            "params": additional_hparams
+        }})
+
+        self._tracker = self._build_tracker()
+
+    def _build_tracker(self) -> BaseTracker:
+        if self._dist_context.is_main_process:
+            return tracker_from_config(self._config.tracker)
+        else:
+            return tracker_from_config(NullTrackerConfig())
+
+    @contextmanager
+    def new_run(self) -> Generator[BaseTrackerRun, None, None]:
+        with self._tracker.open(self._run_config) as run:
+            yield run
+
+    def trigger_sync(self):
+        """
+        Conditionally initiates the synchronization of distributed metrics.
+
+        Checks if the current step is scheduled for metric logging. If so, it
+        triggers the asynchronous communication required to aggregate metric values
+        across ranks. This allows communication to overlap with other operations
+        before `log` is called.
+        """
+
+        if not self._stepper.should_do_action(self._config.period_steps, enable_on_last_step_if_periodic=True):
+            return
+
+        self._metrics.trigger_sync(self._dist_context)
+
+    def log(self, run: BaseTrackerRun, loss_value: torch.Tensor):
+        """
+        Logs the current loss and conditionally processes aggregated metrics.
+
+        This method always logs the provided loss value. Periodically (determined
+        by the stepper and configuration), it waits for the synchronization of
+        metrics to complete (initiated by `trigger_sync`), computes their values,
+        flattens the result structure, logs them to the tracker, and resets the
+        metrics for the next window.
+
+        Args:
+            run: The active tracker run interface for sending data.
+            loss_value: Tensor containing the scalar loss for the current step.
+        """
+
+        run.scalar("loss", loss_value.item())
+
+        if not self._stepper.should_do_action(self._config.period_steps, enable_on_last_step_if_periodic=True):
+            return
+
+        self._metrics.wait_sync(self._dist_context)
+
+        results_tree = self._metrics.compute()
+        results_tree = pytree.tree_map(lambda x: x.item(), results_tree)
+        results_flat = _flatten_pytree_for_metrics(results_tree)
+
+        for name, value in results_flat.items():
+            run.scalar(name, value)
+
+        self._metrics.reset()
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "tracker": state_dict_main_process(self._dist_context, self._tracker),
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        load_state_dict_main_process(self._dist_context, self._tracker, state_dict["tracker"])

d9d/loop/component/job_profiler.py

@@ -0,0 +1,62 @@
+from collections.abc import Generator
+from contextlib import contextmanager
+
+import torch.profiler
+
+from d9d.core.dist_context import DistributedContext
+from d9d.internals.profiling import Profiler
+from d9d.loop.config import ProfilingConfig
+
+from .stepper import Stepper
+
+
+class JobProfiler:
+    """
+    Manages profiling sessions during a job loop.
+
+    This class coordinates the initialization and activation of the internal
+    profiler based on the current step count provided by the stepper.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            config: ProfilingConfig | None,
+            stepper: Stepper
+    ):
+        """
+        Constructs JobProfiler object.
+
+        Args:
+            dist_context: The distributed context.
+            config: Configuration settings for profiling.
+            stepper: Object tracking the current global step of the training loop.
+        """
+
+        self._config = config
+        if config is None or not config.enabled:
+            self._profiler = None
+        else:
+            self._profiler = Profiler(
+                save_dir=config.traces_dir,
+                active_steps=config.active_steps,
+                warmup_steps=config.warmup_steps,
+                period_steps=config.period_steps,
+                dist_context=dist_context
+            )
+        self._stepper = stepper
+
+    @contextmanager
+    def open(self) -> Generator[torch.profiler.profile | None]:
+        """
+        Context manager to activate profiling for the job loop.
+
+        Yields:
+            The active Profiler instance if profiling is enabled, otherwise None.
+        """
+
+        if self._profiler is None:
+            yield None
+        else:
+            with self._profiler.open(self._stepper.current_step) as prof:
+                yield prof

d9d/loop/component/loss_computer.py

@@ -0,0 +1,86 @@
+import torch
+
+from d9d.internals.pipeline_state import PipelineStateHandler
+from d9d.loop.control import ComputeLossContext, TrainTask
+
+from .stepper import Stepper
+
+STATE_LOSS = "__internal_loss"
+STATE_LOSS_WEIGHT = "__internal_loss_weight"
+
+
+class LossComputer:
+    """
+    Handles the computation of loss values and their integration into the pipeline state.
+
+    This component acts as a bridge between the raw outputs of the model pipeline
+    and the user-defined training task. It retrieves the appropriate state context
+    (potentially sharded per microbatch), executes the user's loss logic, persists
+    metrics into the state for logging, and returns the loss*weight term for backpropagation.
+    """
+
+    def __init__(
+            self,
+            state: PipelineStateHandler,
+            task: TrainTask,
+            stepper: Stepper
+    ):
+        """
+        Constructs a new LossComputer.
+
+        Args:
+            state: Handler for managing global and sharded pipeline states.
+            task: The user-defined training task containing loss computation logic.
+            stepper: Component tracking current step and progress.
+        """
+
+        self._state = state
+        self._task = task
+        self._stepper = stepper
+
+    def compute_loss_mul_weight(
+            self,
+            pipeline_outputs: dict[str, torch.Tensor],
+            microbatch_idx: int | None
+    ) -> torch.Tensor:
+        """
+        Computes the weighted loss for a specific sharded microbatch or the full microbatch.
+
+        This method retrieves the appropriate state context based on the microbatch
+        index, delegates calculation to the training task, saves the raw loss and
+        weight into the state for later retrieval, and returns the final scalar
+        product used for backward passes.
+
+        You can retrieve states by using `STATE_LOSS` and `STATE_LOSS_WEIGHT` keys.
+
+        Args:
+            pipeline_outputs: Dictionary containing model output tensors.
+            microbatch_idx: Index of the current microbatch, or `None` for full microbatch execution.
+
+        Returns:
+            The calculated loss multiplied by its weight.
+        """
+
+        if microbatch_idx is None:
+            state = self._state.global_state()
+        else:
+            state = self._state.sharded_state(
+                shard_id=microbatch_idx
+            )
+
+        computation = self._task.compute_loss(ComputeLossContext(
+            pipeline_results=pipeline_outputs,
+            state=state,
+            stepper=self._stepper
+        ))
+
+        loss = computation.loss
+        loss_weight = computation.loss_weight
+
+        if loss_weight is None:
+            loss_weight = torch.ones_like(loss)
+
+        state[STATE_LOSS] = loss[None]
+        state[STATE_LOSS_WEIGHT] = loss_weight[None]
+
+        return loss * loss_weight

d9d/loop/component/model_stage_exporter.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+
+from d9d.core.dist_context import REGULAR_DOMAIN, DistributedContext
+from d9d.loop.control import ModelProvider, PrepareExportModelStageContext
+from d9d.model_state.io import save_model_state_pipeline_parallel
+from d9d.model_state.mapper.compose import ModelStateMapperParallel
+
+from .model_stage_factory import TrackedModules
+
+
+class ModelStageExporter:
+    def __init__(
+            self,
+            model_provider: ModelProvider,
+            modules: TrackedModules,
+            dist_context: DistributedContext
+    ):
+        self._model_provider = model_provider
+        self._modules = modules
+        self._dist_context = dist_context
+
+    def export(self, save_dir: Path):
+        mappers = []
+        for stage in self._modules.modules:
+            result = self._model_provider.prepare_export_model_stage(PrepareExportModelStageContext(
+                model=stage,
+                dist_context=self._dist_context
+            ))
+            mappers.append(result.state_mapper)
+        save_model_state_pipeline_parallel(
+            dest_dir=save_dir,
+            mapper=ModelStateMapperParallel(mappers),
+            device_mesh=self._dist_context.mesh_for(REGULAR_DOMAIN),
+            pipeline_dim_name="pp",
+            models=self._modules.modules,
+            show_progress=True,
+        )

d9d/loop/component/model_stage_factory.py

@@ -0,0 +1,261 @@
+import itertools
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch import nn
+from torch.distributed.checkpoint.stateful import Stateful
+
+from d9d.core.dist_context import REGULAR_DOMAIN, DistributedContext
+from d9d.loop.config import ModelStageFactoryConfig, PipeliningConfig
+from d9d.loop.control import InitializeModelStageContext, ModelProvider, ParallelizeModelStageContext
+from d9d.model_state.io import load_model_state
+from d9d.module.base import ModuleLateInit
+from d9d.pipelining.api import PipelineStageInfo
+from d9d.pipelining.factory.factory import PipelineScheduleInfo, build_schedule
+
+from .batch_maths import BatchMaths
+from .loss_computer import LossComputer
+
+StatefulPredicate = Callable[[str, torch.Tensor], bool]
+"""Determines if a specific parameter or buffer should be included in the state dictionary."""
+
+
+def _stateful_predicate_requires_grad(key: str, value: torch.Tensor) -> bool:
+    """Predicate that allows saving only tensors that require gradients."""
+    return value.requires_grad
+
+
+def _stateful_predicate_always(key: str, value: torch.Tensor) -> bool:
+    """Predicate that always allows saving."""
+    return True
+
+
+class TrackedModules(Stateful):
+    """
+    Wraps a list of model stages and manages their state for distributed checkpointing.
+
+    This class implements the PyTorch Distributed `Stateful` protocol, aggregating
+    the state dictionaries of multiple pipeline stages assigned to the current rank.
+    It handles namespacing to ensure uniqueness across pipeline ranks and stages.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            modules: list[nn.Module],
+            stateful_predicate: StatefulPredicate
+    ):
+        """Constructs a TrackedModules object."""
+        self._dist_context = dist_context
+        self._modules = modules
+        self._stateful_predicate = stateful_predicate
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Forwards execution to the only pipeline stage.
+
+        This method is only valid when pipeline parallelism is disabled.
+
+        Args:
+            *args: Positional arguments passed to the module.
+            **kwargs: Keyword arguments passed to the module.
+
+        Returns:
+            The output of the model execution.
+
+        Raises:
+            ValueError: If pipeline parallelism is configured.
+        """
+
+        if self._dist_context.mesh_params.has_pipeline_parallel:
+            raise ValueError("You cannot call tracked modules when using pipelining")
+
+        return self._modules[0](*args, **kwargs)
+
+    @property
+    def modules(self) -> list[nn.Module]:
+        """Returns the list of underlying PyTorch model modules."""
+        return self._modules
+
+    def _whitelisted_params(self, module: nn.Module) -> set[str]:
+        allow_saving = set()
+        for param_name, param in itertools.chain(module.named_parameters(), module.named_buffers()):
+            if self._stateful_predicate(param_name, param):
+                allow_saving.add(param_name)
+        return allow_saving
+
+    def _state_dict_stage(self, module: nn.Module) -> dict[str, Any]:
+        whitelist = self._whitelisted_params(module)
+        result = {
+            k: v for k, v in module.state_dict().items() if k in whitelist
+        }
+        return result
+
+    def state_dict(self) -> dict[str, Any]:
+        """
+        Generates the state dictionary for all tracked modules.
+
+        The keys are namespaced using the current pipeline rank and stage index
+        (e.g., `pp_0_stage_0`). Only parameters satisfying the `stateful_predicate`
+        are included.
+
+        Returns:
+            A dictionary containing the states of all managed modules.
+        """
+
+        pp_rank = self._dist_context.mesh_for(REGULAR_DOMAIN)["pp"].get_local_rank()
+        ret = {
+            f"pp_{pp_rank}_stage_{i}": self._state_dict_stage(module)
+            for i, module in enumerate(self._modules)
+        }
+        return ret
+
+    def _load_state_dict_stage(self, module: nn.Module, state_dict: dict[str, Any]):
+        whitelist = self._whitelisted_params(module)
+
+        loading_result = module.load_state_dict(state_dict, strict=False)
+        missing_keys = set(loading_result.missing_keys)
+        extra_keys = set(loading_result.unexpected_keys)
+
+        if len(whitelist.intersection(missing_keys)) > 0:
+            raise ValueError(f"Missing keys: {whitelist.intersection(missing_keys)}")
+        if len(extra_keys) > 0:
+            raise ValueError(f"Extra keys: {extra_keys}")
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """
+        Loads the state dictionary into the tracked modules.
+
+        Args:
+            state_dict: The state dictionary to load. Must contain keys corresponding
+                to the pipeline rank and stage indices managed by this instance.
+
+        Raises:
+            ValueError: If required keys are missing or unexpected keys are present
+                based on the allow-list predicate.
+        """
+
+        pp_rank = self._dist_context.mesh_for(REGULAR_DOMAIN)["pp"].get_local_rank()
+        for i, module in enumerate(self._modules):
+            self._load_state_dict_stage(module, state_dict[f"pp_{pp_rank}_stage_{i}"])
+
+
+class ModelStageFactory:
+    """
+    Factory class responsible for creating, initializing, and parallelizing model stages.
+
+    This class coordinates the `ModelProvider` with the distributed context to:
+
+    1. Initialize models on a meta device.
+    2. Apply horizontal distribution strategy (TP, DP, FSDP, etc).
+    3. Materialize weights on the target device.
+    4. Load initial model states from checkpoints.
+    """
+
+    def __init__(
+            self,
+            model_provider: ModelProvider,
+            dist_context: DistributedContext,
+            batch_maths: BatchMaths,
+            config_model: ModelStageFactoryConfig,
+            config_pipelining: PipeliningConfig | None,
+            loss_computer: LossComputer | None
+    ):
+        """Constructs a ModelStageFactory object."""
+
+        self._model_provider = model_provider
+        self._dist_context = dist_context
+        self._config_model = config_model
+        self._config_pipelining = config_pipelining
+        self._batch_maths = batch_maths
+        self._loss_computer = loss_computer
+
+    def _build_model_stage(self, stage: PipelineStageInfo) -> nn.Module:
+        # create a model with no real memory occupied
+        with torch.device("meta"):
+            factored = self._model_provider.initialize_model_stage(
+                InitializeModelStageContext(
+                    dist_context=self._dist_context,
+                    stage=stage,
+                )
+            )
+
+        model = factored.model
+
+        if not isinstance(model, ModuleLateInit) or not isinstance(model, nn.Module):
+            raise ValueError("Model stage is required to be nn.Module instance implementing ModuleLateInit protocol")
+
+        # if current context is distributed - parallelize this model
+        if self._dist_context.mesh_params.is_distributed:
+            self._model_provider.parallelize_model_stage(
+                ParallelizeModelStageContext(
+                    model=model,
+                    stage=stage,
+                    dist_context=self._dist_context
+                )
+            )
+
+        # move state that is bound to current device to it
+        model.to_empty(device=self._dist_context.current_device)
+
+        # reinitialize model parameters (only these are on current device)
+        with torch.no_grad():
+            model.reset_parameters()
+
+        if self._config_model.source_checkpoint:
+            load_model_state(
+                src_dir=self._config_model.source_checkpoint,
+                model=model,
+                mapper=factored.state_mapper,
+                device=f"cuda:{torch.cuda.current_device()}"
+            )
+
+        # set training state
+        model.train()
+
+        return model
+
+    def build_pipeline_and_modules(
+            self
+    ) -> tuple[PipelineScheduleInfo | None, TrackedModules]:
+        """
+        Constructs the execution schedule and the model container.
+
+        If pipeline parallelism is enabled, this orchestrates the creation of a
+        distributed pipeline schedule.
+
+        Otherwise, it simply builds a standalone model stage.
+
+        Returns:
+           The pipeline schedule information (or None if no pipelining).
+           The `TrackedModules` instance wrapping the created model stage(s).
+
+        Raises:
+           ValueError: If pipelining configuration is missing but a pipeline is requested.
+        """
+
+        if self._config_model.checkpoint_only_trainable_parameters:
+            stateful_predicate = _stateful_predicate_requires_grad
+        else:
+            stateful_predicate = _stateful_predicate_always
+
+        if self._dist_context.mesh_params.has_pipeline_parallel:
+            if self._config_pipelining is None:
+                raise ValueError("Pipelining is enabled, but not configured")
+
+            loss_fn = self._loss_computer.compute_loss_mul_weight if self._loss_computer is not None else None
+
+            schedule, modules = build_schedule(
+                dist_context=self._dist_context,
+                n_microbatches=self._batch_maths.num_microbatches_pipelining,
+                schedule_config=self._config_pipelining.schedule,
+                model_provider=self._build_model_stage,
+                loss_fn=loss_fn
+            )
+
+            return schedule, TrackedModules(self._dist_context, modules, stateful_predicate)
+        else:
+            model = self._build_model_stage(PipelineStageInfo(num_stages=1, current_stage=0))
+
+            return None, TrackedModules(self._dist_context, [model], stateful_predicate)

d9d/loop/component/optimizer_factory.py

@@ -0,0 +1,88 @@
+from d9d.core.dist_context import REGULAR_DOMAIN, DistributedContext
+from d9d.core.protocol import LRSchedulerProtocol, OptimizerProtocol
+from d9d.loop.control import (
+    InitializeLRSchedulerContext,
+    InitializeOptimizerStageContext,
+    LRSchedulerProvider,
+    OptimizerProvider,
+)
+from d9d.pipelining.training import PipelinedLRScheduler, PipelinedOptimizer
+
+from .model_stage_factory import TrackedModules
+from .stepper import Stepper
+
+
+class OptimizerFactory:
+    """
+    Factory for creating and configuring distributed optimizers and learning rate schedulers.
+
+    This factory handles the orchestration of optimizer creation for models potentially split across
+    pipeline stages. It uses the providers to instantiate underlying PyTorch optimizers and schedulers for each
+    tracked module, and wraps them in pipeline-aware interfaces.
+    """
+
+    def __init__(
+            self,
+            dist_context: DistributedContext,
+            tracked_modules: TrackedModules,
+            optimizer_provider: OptimizerProvider,
+            lr_scheduler_provider: LRSchedulerProvider,
+            stepper: Stepper
+    ):
+        """
+        Constructs the OptimizerFactory.
+
+        Args:
+            dist_context: The distributed context.
+            tracked_modules: A container of model modules owned by the current rank.
+            optimizer_provider: A callable responsible for creating optimizer instances for a given model.
+            lr_scheduler_provider: A callable responsible for creating LR scheduler instances.
+            stepper: The training stepper providing information about total training steps.
+        """
+        self._dist_context = dist_context
+        self._tracked_modules = tracked_modules
+        self._optimizer_provider = optimizer_provider
+        self._lr_scheduler_provider = lr_scheduler_provider
+        self._stepper = stepper
+
+    def build_optimizer_and_scheduler(self) -> tuple[OptimizerProtocol, LRSchedulerProtocol]:
+        """
+        Builds both the optimizer and learning rate scheduler.
+
+        This method iterates through all local model modules. For each module, it creates an
+        optimizer and scheduler using the configured providers. Finally, it aggregates these individual
+        instances into a single `PipelinedOptimizer` and `PipelinedLRScheduler` capable of coordinated
+        stepping across the pipeline parallel dimension.
+
+        Returns:
+            A tuple containing the initialized pipeline-aware optimizer and scheduler.
+        """
+
+        optimizers: list[OptimizerProtocol] = []
+        lr_schedulers: list[LRSchedulerProtocol] = []
+        for module in self._tracked_modules.modules:
+            optimizer = self._optimizer_provider(
+                InitializeOptimizerStageContext(
+                    dist_context=self._dist_context,
+                    model=module
+                )
+            )
+            optimizers.append(optimizer)
+
+            scheduler = self._lr_scheduler_provider(
+                InitializeLRSchedulerContext(
+                    dist_context=self._dist_context,
+                    total_steps=self._stepper.total_steps,
+                    optimizer=optimizer
+                )
+            )
+            lr_schedulers.append(scheduler)
+        pipe_optimizer = PipelinedOptimizer(
+            mesh_pp=self._dist_context.mesh_for(REGULAR_DOMAIN)["pp"],
+            optimizers=optimizers
+        )
+        pipe_scheduler = PipelinedLRScheduler(
+            mesh_pp=self._dist_context.mesh_for(REGULAR_DOMAIN)["pp"],
+            schedulers=lr_schedulers
+        )
+        return pipe_optimizer, pipe_scheduler