TorchForge

https://meta-pytorch.org/torchforge
Admin
TorchForge is a PyTorch-native library for reinforcement learning post-training and agentic...
Tokens:73,835
Snippets:495
Trust Score:6.4
Update:1 month ago
Show doc for...
Tokens:
Raw
### Run TorchForge installation script

Source: https://meta-pytorch.org/torchforge/main/getting_started.html

Runs the installation script that installs system dependencies, downloads pre-built wheels, and configures the environment for GPU training. Use --use-sudo to install system packages directly instead of via conda.

```bash
./scripts/install.sh

```

--------------------------------

### ForgeActor.options usage examples

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/actor.html

Four usage examples from the `options` method docstring. The first two show pre-configuring a service with `num_replicas=2, procs=2` and the default service usage. The last two show pre-configuring a single actor with `procs=1, hosts=1` and the default actor usage. Each example calls `shutdown()` after spawning.

```python
service = await MyForgeActor.options(num_replicas=2, procs=2).as_service(...)
await service.shutdown()
```

```python
service = await MyForgeActor.as_service(...)
await service.shutdown()
```

```python
actor = await MyForgeActor.options(procs=1, hosts=1).as_actor(...)
await actor.shutdown()
```

```python
actor = await MyForgeActor.as_actor(...)
await MyForgeActor.shutdown(actor)
```

--------------------------------

### Install TorchForge and Run Training

Source: https://meta-pytorch.org/torchforge/main/getting_started.html

Commands to set up the environment, clone the repository, install dependencies, download a model, and run SFT and GRPO training. Requires 2+ GPUs for training.

```bash
# Install dependencies
conda create -n forge python=3.10
conda activate forge
git clone https://github.com/meta-pytorch/forge
cd forge
./scripts/install.sh

# Download a model
hf download meta-llama/Meta-Llama-3.1-8B-Instruct --local-dir /tmp/Meta-Llama-3.1-8B-Instruct --exclude "original/consolidated.00.pth"

# Run SFT training (requires 2+ GPUs)
uv run forge run --nproc_per_node 2 \
  apps/sft/main.py --config apps/sft/llama3_8b.yaml

# Run GRPO training (requires 2+ GPUs)
python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml
```

--------------------------------

### Initialize replicas and start health loop

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Initializes replicas in parallel and starts the background health loop. Logs the number of queued replicas before initialization.

```python
logger.debug(
    f"Queued {num_replicas} replicas for initialization. Total replicas: {len(self._replicas)}"
)

# Initialize all replicas in parallel
await asyncio.gather(*[r.initialize() for r in replicas])
self._replicas = replicas

# Start the health loop in the background
self._health_task = asyncio.create_task(
    self._health_loop(poll_rate_s=self._cfg.health_poll_rate)
)
```

--------------------------------

### Service.start_session()

Source: https://meta-pytorch.org/torchforge/main/api_service.html

Starts a new session for stateful request handling. Sessions enable request affinity to specific replicas, maintaining state consistency for workloads that require it. Each session gets a unique ID and is automatically assigned to the least loaded replica.

```APIDOC
## Service.start_session()

### Description
Starts a new session for stateful request handling. Sessions enable request affinity to specific replicas, maintaining state consistency for workloads that require it. Each session gets a unique ID and is automatically assigned to the least loaded replica.

### Method
start_session

### Returns
- **str** - Unique session identifier for use in subsequent requests

### Example
```
>>> session_id = await service.start_session()
>>> result = await service.my_endpoint(session_id, arg1, arg2)
>>> await service.terminate_session(session_id)
```
```

--------------------------------

### Generator.launch - provision host mesh, allocate GPUs, spawn procs

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/actors/vllm/v1/generator.html

This is the launch method of the Generator actor. It retrieves the provisioner, gets a host mesh (remote or local), mounts wsfuse on remote hosts, allocates GPU IDs, spawns a CPU proc for the Generator and WorkerRegistry, and finally calls setup() with the host mesh and GPU IDs. It also registers a metric logger if metrics are enabled.

```python
        engine_args = kwargs.get("engine_args", {})
        if isinstance(engine_args, Mapping):
            engine_args = EngineArgs(**engine_args)
        vllm_config = engine_args.create_engine_config(UsageContext.LLM_CLASS)

        num_gpus = vllm_config.parallel_config.world_size
        num_hosts = cls.hosts if cls.hosts else None
        gpus_per_host = num_gpus // (num_hosts or 1)
        mesh_name = cls.mesh_name or "generator"

        # Step 1: Get host mesh directly (no bootstrap procs)
        provisioner = await _get_provisioner()
        if provisioner is None:
            raise RuntimeError(
                "Provisioner not initialized. Call init_provisioner() first."
            )

        if num_hosts:
            # Remote allocation - get pre-allocated mesh from launcher
            host_mesh = await provisioner.get_host_mesh(mesh_name)
        else:
            # Local allocation
            host_mesh = this_host()

        # Step 1a: Mount wsfuse on all hosts for local model paths
        # This must happen before any actors are spawned that might access the path
        if num_hosts and provisioner.launcher:
            # Spawn temporary procs on all hosts just for mounting
            mount_procs = host_mesh.spawn_procs(
                per_host={"procs": 1}, name="mount_setup"
            )
            await provisioner.launcher.remote_setup(mount_procs)
            logger.info("Completed remote_setup (mounted wsfuse on all hosts)")

        # Step 1b: Allocate GPUs from provisioner
        # This ensures each Generator gets exclusive GPU IDs
        gpu_ids = await provisioner.allocate_gpu_ids(host_mesh, num_gpus)
        logger.info(f"[Generator.launch] Allocated GPUs: {gpu_ids}")

        logger.info(
            f"[Generator.launch] Provisioned host mesh with {num_hosts or 1} host(s), "
            f"{gpus_per_host} GPUs per host"
        )

        # Step 2: Spawn CPU proc on head host for Generator and WorkerRegistry
        if num_hosts:
            singleton_slice = {k: slice(0, 1) for k in host_mesh.extent.keys()}
            head_host = host_mesh.slice(**singleton_slice)
        else:
            head_host = host_mesh

        generator_proc = head_host.spawn_procs(
            per_host={"procs": 1}, name="generator_proc"
        )
        logger.info("[Generator.launch] Spawned generator_proc on head host")

        # Register LocalFetcherActor for generator_proc to enable metrics collection
        if not FORGE_DISABLE_METRICS.get_value():
            await get_or_create_metric_logger(generator_proc, process_name=mesh_name)

        # Import WorkerRegistry here to avoid circular import with monarch_executor
        from forge.actors.vllm.v1.monarch_executor import WorkerRegistry

        # Spawn WorkerRegistry on CPU proc (same as Generator)
        worker_registry = generator_proc.spawn("worker_registry", WorkerRegistry)
        logger.info("[Generator.launch] Spawned WorkerRegistry on generator_proc")

        actor_name = kwargs.pop("name", cls.__name__)
        generator = generator_proc.spawn(
            actor_name,
            cls,
            *args,
            **kwargs,
        )

        # Attach for cleanup in Generator.shutdown()
        generator._generator_proc = generator_proc
        generator._worker_registry = worker_registry

        # Step 3: Pass host_mesh and gpu_ids to setup() - executor will create proc_mesh
        await generator.setup.call(host_mesh, worker_registry, gpu_ids)
        return generator
```

--------------------------------

### Start, use, and terminate a session

Source: https://meta-pytorch.org/torchforge/main/api_service.html

Starts a session, uses it for a request, and terminates it. Sessions enable request affinity to a specific replica and should be terminated when no longer needed.

```python
>>> session_id = await service.start_session()
>>> result = await service.my_endpoint(session_id, arg1, arg2)
>>> await service.terminate_session(session_id)
```

--------------------------------

### Tracer.start()

Source: https://meta-pytorch.org/torchforge/main/metric_logging.html

Starts the tracer. This is called on a Tracer instance.

```APIDOC
## Tracer.start()

### Description
Starts the tracer. This is called on a Tracer instance.

### Method
Method call

### Endpoint
`t.start()`

### Parameters
None.

### Request Example
```python
t.start()
```

### Response
No return value.

```

--------------------------------

### Initialize service and start health loop

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Endpoint method that initializes the service by creating the specified number of replicas with the given configuration and actor definition. Logs the number of replicas being started.

```python
@endpoint
async def __initialize__(self):
    """Initializes the service and starts the health loop."""
    logger.debug(f"Starting service up with {self._cfg.num_replicas} replicas.")
    replicas = []
    num_replicas = self._cfg.num_replicas
    for i in range(num_replicas):
        replica = Replica(
            idx=len(self._replicas) + i,
            proc_config=self._cfg.to_process_config(),
            max_concurrent_requests=self._cfg.replica_max_concurrent_requests,
            return_first_rank_result=self._cfg.return_first_rank_result,
            actor_def=self._actor_def,
            actor_kwargs=self._actor_kwargs,
        )
        replicas.append(replica)
```

--------------------------------

### Configure multiple backends with init_backends.call_one

Source: https://meta-pytorch.org/torchforge/main/metric_logging.html

Shows how to configure multiple logging backends at once using init_backends.call_one. The example sets console to per_rank_no_reduce and wandb to global_reduce, allowing different logging modes per backend.

```python
mlogger = await get_or_create_metric_logger(process_name="Controller")
await mlogger.init_backends.call_one({
    "console": {"logging_mode": "per_rank_no_reduce"},
    "wandb": {"logging_mode": "global_reduce"}
})
```

--------------------------------

### Generator usage example

Source: https://meta-pytorch.org/torchforge/main/api_generator.html

Shows how to create a Generator service, generate a completion, and shut it down. Uses EngineArgs and SamplingParams as dictionaries.

```python
>>> generator = await Generator.options(procs=1, with_gpus=True).as_service(
...     engine_args={"model": "meta-llama/Llama-3-8B", "tensor_parallel_size": 2},
...     sampling_params={"max_tokens": 128, "temperature": 0.7},
... )
>>> completions = await generator.generate("Tell me a joke")
>>> await generator.shutdown()
```

--------------------------------

### Asynchronous RL Step with TorchForge Services and Actors

Source: https://meta-pytorch.org/torchforge/main/tutorials/zero-to-forge/2_Forge_Internals.html

Shows a complete RL step using TorchForge's service abstraction. It demonstrates calling actors (dataloader, compute_advantages, replay_buffer) and services (generator, ref_model, reward_actor, trainer) with .call_one() and .route() methods. The example creates an Episode, computes advantages, stores in replay buffer, and trains on a batch. Note the comment about tokenizer and services being pre-initialized.

```python
from apps.grpo.main import Episode, Group

async def simple_rl_step():

    # ===== Generate a rollout =====
    sample = await dataloader.sample.call_one()  # DatasetActor is an actor, not service
    prompt, target = sample["request"], sample["target"]  # Correct field names

    print(f"Prompt: {prompt}")
    print(f"Target: {target}")

    actions = await generator.generate.route(prompt=prompt)  # Generator is a service
    print(f"Generator response: {actions[0].text}")

    # Create input tensor for reference model (requires full context)
    input_ids = torch.cat([actions[0].prompt_ids, actions[0].token_ids])
    ref_logprobs = await ref_model.forward.route(
        input_ids.unsqueeze(0), max_req_tokens=512, return_logprobs=True
    )
    reward = await reward_actor.evaluate_response.route(  # RewardActor is a service
        prompt=prompt,
        response=actions[0].text,
        target=target
    )
    print(f"Reward: {reward}")

    # Create episode using actual GRPO Episode structure
    episode = Episode(
        episode_id="0",
        request=prompt,
        policy_version=0,
        pad_id=tokenizer.pad_token_id,
        request_len=512,
        response_len=512,
        target=target
    )

    # Add response data
    episode.response = actions[0].text
    episode.request_tokens = actions[0].prompt_ids.tolist()
    episode.response_tokens = actions[0].token_ids.tolist()
    episode.ref_logprobs = ref_logprobs[0]  # Extract from batch dimension
    episode.reward = reward

    # Compute advantages using actual ComputeAdvantages actor
    group = Group.new_group(0, 1, prompt, 0, tokenizer.pad_token_id, 512, 512, target)
    group.episodes[0] = episode
    advantages = await compute_advantages.compute.call_one(group)  # ComputeAdvantages is an actor
    episode.advantage = advantages[0]
    print(f"Advantage: {advantages[0]}")
    await replay_buffer.add.call_one(episode)  # ReplayBuffer is an actor
    print("Episode stored in replay buffer")

    # ===== Train on the batch =====
    batch = await replay_buffer.sample.call_one(curr_policy_version=0)
    if batch is not None:
        print("Training on batch...")
        inputs, targets = batch  # GRPO returns (inputs, targets) tuple
        loss = await trainer.train_step.call(inputs, targets)  # TitanTrainer is an actor
        print(f"Training loss: {loss}")
        return loss
    else:
        print("Not enough data in buffer yet")
        return None

# Note: This simplified example assumes tokenizer and services are already initialized
for step in range(10):
    print(f"\n--- RL Step {step + 1} ---")
    loss = await simple_rl_step()
    if loss:
        print(f"Step {step + 1} complete, loss: {loss:.4f}")
    else:
        print(f"Step {step + 1} complete, building buffer...")
```

--------------------------------

### Custom backend example: ConsoleBackend extending LoggerBackend

Source: https://meta-pytorch.org/torchforge/main/metric_logging.html

Example of a custom backend that extends LoggerBackend. It implements init, log_batch (called on flush), and log_stream (called on record_metric when logging_mode is per_rank_no_reduce). Register with forge.observability.metrics.get_logger_backend_class.

```python
# Example of a custom backend
class ConsoleBackend(LoggerBackend):
    def __init__(self, logger_backend_config: dict[str, Any]) -> None:
        super().__init__(logger_backend_config)

    async def init(self, process_name: str | None = None, *args, **kwargs) -> None:
        self.process_name = process_name

    async def log_batch(self, metrics: list[Metric], global_step: int, *args, **kwargs) -> None:
        # Called on flush
        print(self.process_name, metrics)

    def log_stream(self, metric: Metric, global_step: int, *args, **kwargs) -> None:
        # Called on `record_metric` if "logging_mode": "per_rank_no_reduce"
        print(metric)
```

--------------------------------

### TorchForge YAML training configuration

Source: https://meta-pytorch.org/torchforge/main/getting_started.html

Example YAML configuration file for TorchForge training, showing model, training, distributed, and checkpointing sections. The distributed strategy is managed by TorchTitan.

```yaml
# Example: apps/sft/llama3_8b.yaml
model:
  name: meta-llama/Meta-Llama-3.1-8B-Instruct
  path: /tmp/Meta-Llama-3.1-8B-Instruct

training:
  batch_size: 4
  learning_rate: 1e-5
  num_epochs: 10
  gradient_accumulation_steps: 4

distributed:
  strategy: fsdp  # Managed by TorchTitan
  precision: bf16

checkpointing:
  save_interval: 1000
  output_dir: /tmp/checkpoints

```

--------------------------------

### Generator.setup - initialize AsyncLLM with MonarchExecutor

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/actors/vllm/v1/generator.html

This is the setup method of the Generator actor. It receives a host_mesh, worker_registry, and GPU IDs from launch(). It sets the VLLM_MONARCH_GPU_IDS environment variable for the MonarchExecutor subprocess (EngineCore). The MonarchExecutor creates its own proc_mesh from host_mesh, spawns workers, and registers them.

```python
        num_gpus = self.vllm_config.parallel_config.tensor_parallel_size

        logger.info(f"Setting up AsyncLLM with {num_gpus} GPUs, allocated: {gpu_ids}")

        # Set env var for MonarchExecutor subprocess (EngineCore)
        os.environ["VLLM_MONARCH_GPU_IDS"] = ",".join(gpu_ids)
```

--------------------------------

### Implement RewardActor service with evaluate_response endpoint

Source: https://meta-pytorch.org/torchforge/main/tutorials/zero-to-forge/2_Forge_Internals.html

Defines a RewardActor class that extends ForgeActor, with an endpoint method evaluate_response that averages the scores from multiple reward functions. The service is started with RewardActor.options(...).as_service(...) and called via the route method. Includes a sample prompt/response/target and prints the reward score. Also shows how to scale with num_replicas and cleanup with shutdown().

```python
# Exact RewardActor from apps/grpo/main.py

from forge.controller import ForgeActor
from monarch.actor import endpoint
from forge.data.rewards import MathReward, ThinkingReward

# class definition from apps/grpo/main.py
class RewardActor(ForgeActor):
    def __init__(self, reward_functions: list):
        self.reward_functions = reward_functions

    @endpoint
    async def evaluate_response(self, prompt: str, response: str, target: str) -> float:
        """Evaluate response quality using multiple reward functions"""
        total_reward = 0.0

        for reward_fn in self.reward_functions:
            # Each reward function contributes to total score
            reward = reward_fn(prompt, response, target)
            total_reward += reward

        # Return average reward across all functions
        return total_reward / len(self.reward_functions) if self.reward_functions else 0.0

reward_actor = await RewardActor.options(
    procs=1, num_replicas=1
).as_service(
    reward_functions=[MathReward(), ThinkingReward()]
)

prompt = "What is 15% of 240?"
response = "15% of 240 is 36"
target = "36"

score = await reward_actor.evaluate_response.route(
    prompt=prompt,
    response=response,
    target=target
)
print(f"Reward score: {score}")  # Usually around 1.0 for correct math answers

# For production scaling - increase num_replicas for parallel evaluation:
# RewardActor.options(procs=1, num_replicas=16)  # 16 parallel evaluators

# Cleanup when done
await reward_actor.shutdown()

```

--------------------------------

### ForgeActor service sessions for sticky routing

Source: https://meta-pytorch.org/torchforge/main/tutorials/zero-to-forge/2_Forge_Internals.html

Shows how to use ForgeActor sessions to route multiple calls to the same replica, contrasting with regular .route() calls that load-balance across replicas. The first block defines a ForgeCounter service and demonstrates inconsistent state without sessions; the second block uses counter_service.session() to ensure all calls hit the same replica, preserving state and KV cache. The commented example shows the same pattern with a Policy for multi-turn conversations.

```python
# This Counter example demonstrates the difference between regular routing and sessions

from forge.controller import ForgeActor
from monarch.actor import endpoint

class ForgeCounter(ForgeActor):
    def __init__(self, initial_value: int):
        self.value = initial_value

    @endpoint
    def increment(self) -> int:
        self.value += 1
        return self.value

    @endpoint
    def get_value(self) -> int:
        return self.value

    @endpoint
    async def reset(self):
        self.value = 0

counter_service = await ForgeCounter.options(
    procs=1, num_replicas=4
).as_service(initial_value=0)

# WITHOUT SESSIONS: Each .route() call goes to a different replica
await counter_service.increment.route()  # Might go to replica 2
await counter_service.increment.route()  # Might go to replica 1
await counter_service.increment.route()  # Might go to replica 3

results = await counter_service.increment.fanout()  # Get from all replicas
print(f"All replica values: {results}")
# Output: All replica values: [1, 2, 1, 1] - Each replica has different state!
```

```python
# WITH SESSIONS: All calls go to the SAME replica
print("\nUsing sticky sessions:")
async with counter_service.session():  # Creates a session that picks one replica
    await counter_service.reset.route()  # Uses .route() within session
    print(await counter_service.increment.route())  # 1
    print(await counter_service.increment.route())  # 2
    print(await counter_service.increment.route())  # 3

    final_value = await counter_service.get_value.route()
    print(f"Final value on this replica: {final_value}")  # 3

# Output:
# Using sticky sessions:
# 1
# 2
# 3
# Final value on this replica: 3

# Same pattern works with Policy for multi-turn conversations:
# async with generator.session():
#     response1 = await generator.generate.route(turn1)
#     full_prompt = turn1 + response1[0].text + turn2
#     response2 = await generator.generate.route(full_prompt)
#     # Both calls hit same replica, preserving KV cache

# Cleanup
await counter_service.shutdown()
```

--------------------------------

### Start session for stateful request handling

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Starts a new session for stateful request handling. Generates a unique session ID, creates a Session object, appends it to the active sessions list, updates metrics, and returns the session ID. The docstring includes an example usage showing start_session, calling an endpoint, and terminate_session.

```python
    @endpoint
    async def start_session(self) -> str:
        """
        Starts a new session for stateful request handling.

        Sessions enable request affinity to specific replicas, maintaining state
        consistency for workloads that require it. Each session gets a unique ID
        and is automatically assigned to the least loaded replica.

        Returns:
            str: Unique session identifier for use in subsequent requests

        Example:
            >>> session_id = await service.start_session()
            >>> result = await service.my_endpoint(session_id, arg1, arg2)
            >>> await service.terminate_session(session_id)
        """
        sess_id = str(uuid.uuid4())
        session = Session(session_id=sess_id)
        self._active_sessions.append(session)

        # Update metrics
        self._update_service_metrics()

        return sess_id
```

--------------------------------

### Get comprehensive service metrics

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Returns comprehensive service metrics including per-replica performance data, service-wide aggregations, and health status. The docstring shows an example of retrieving metrics and printing request rate and queue depth.

```python
    @endpoint
    def get_metrics(self) -> ServiceMetrics:
        """
        Get comprehensive service metrics for monitoring and analysis.

        Returns detailed metrics including per-replica performance data,
        service-wide aggregations, and health status information.

        Returns:
            ServiceMetrics: Complete metrics object with replica and service data

        Example:
            >>> metrics = service.get_metrics()
            >>> print(f"Request rate: {metrics.get_total_request_rate():.1f} req/s")
            >>> print(f"Queue depth: {metrics.get_avg_queue_depth():.1f}")
        """
        self._update_service_metrics()
        return self._metrics
```

--------------------------------

### Record and log a metric with configurable logging mode

Source: https://meta-pytorch.org/torchforge/main/metric_logging.html

Sets up a ForgeActor that records a metric (my_sum_rank_metric) with Reduce.SUM, initializes a metric logger with a chosen logging_mode (e.g., global_reduce), runs the actor across 2 replicas and 2 processes, and flushes the logger. The example uses current_rank().rank to get the rank value (0 or 1) and records it for 2 steps.

```python
import asyncio

from forge.controller.actor import ForgeActor
from forge.observability import get_or_create_metric_logger, record_metric, Reduce
from monarch.actor import current_rank, endpoint

# Your distributed actor
class MyActor(ForgeActor):
    @endpoint
    async def my_fn(self):
        rank = current_rank().rank # 0 or 1 per replica
        record_metric("my_sum_rank_metric", rank, Reduce.SUM) # <--- your metric

async def main():
    # Setup logger
    mlogger = await get_or_create_metric_logger(process_name="Controller")
    await mlogger.init_backends.call_one(
        {"console": {"logging_mode": "global_reduce"}} #  <--- Define logging_mode here
    )

    # Setup actor
    service_config = {"procs": 2, "num_replicas": 2, "with_gpus": False}
    my_actor = await MyActor.options(**service_config).as_service()

    # Accumulate metrics
    for _ in range(2):  # 2 steps
        await my_actor.my_fn.fanout()

    # Flush
    await mlogger.flush.call_one(global_step=0)  # Flush and reset

if __name__ == "__main__":
    asyncio.run(main())
```

--------------------------------

### Pre-configure a service with multiple replicas using options()

Source: https://meta-pytorch.org/torchforge/main/api_actors.html

Shows how to pre-configure a ForgeActor subclass with num_replicas=2 and procs=2 before spawning as a service. The service is then shut down with await service.shutdown().

```python
service = await MyForgeActor.options(num_replicas=2, procs=2).as_service(...)
await service.shutdown()
```

--------------------------------

### Naive RL Step with Bottlenecks

Source: https://meta-pytorch.org/torchforge/main/tutorials/zero-to-forge/1_RL_and_Forge_Fundamentals.html

Shows a naive RL step that serializes policy generation, reward evaluation, and training, causing GPU idle time, batch size 1, and a halt if any component fails. Contrasts with TorchForge's coordination and failover.

```python
def naive_rl_step():
    # Policy waits idle while reward model works
    response = policy_model.generate(prompt)  # GPU busy
    reward = reward_model.evaluate(prompt, response)  # Policy GPU idle

    # Training waits for single episode
    loss = compute_loss(response, reward)  # Batch size = 1, inefficient

    # Everything stops if any component fails
    if policy_fails or reward_fails or trainer_fails:
        entire_system_stops()
```

--------------------------------

### Custom reduce type example: MeanAccumulator extending MetricAccumulator

Source: https://meta-pytorch.org/torchforge/main/metric_logging.html

Example of a custom reduce type that implements Reduce.MEAN. It accumulates values, computes the mean, and supports global reduce via get_reduced_value_from_states. Register with forge.observability.metrics.Reduce.

```python
# Example of a custom reduce type
class MeanAccumulator(MetricAccumulator):
    def __init__(self, reduction: Reduce) -> None:
        super().__init__(reduction)
        self.sum = 0.0
        self.count = 0
        self.is_reset = True

    def append(self, value: Any) -> None:
        # Called after record_metric(key, value, reduce.TYPE)
        v = float(value.item() if hasattr(value, "item") else value)
        self.sum += v
        self.count += 1

    def get_value(self) -> float:
        return self.sum / self.count if self.count > 0 else 0.0

    def get_state(self) -> dict[str, Any]:
        return {"reduction_type": self.reduction_type.value, "sum": self.sum, "count": self.count}

    @classmethod
    def get_reduced_value_from_states(cls, states: list[dict[str, Any]]) -> float:
        # Useful for global reduce; called before flush
        total_sum = sum(s["sum"] for s in states)
        total_count = sum(s["count"] for s in states)
        return total_sum / total_count if total_count > 0 else 0.0

    def reset(self) -> None:
        self.sum = 0.0
        self.count = 0
        self.is_reset = True
```

--------------------------------

### Terminate session example

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Shows how to terminate an active session using the terminate_session method. The session ID is obtained from start_session.

```python
session_id = await service.start_session()
>>> # ... use session for requests ...
>>> await service.terminate_session(session_id)
```

--------------------------------

### record_metric with Reduce.SUM/MAX/MEAN and console backend output

Source: https://meta-pytorch.org/torchforge/main/metric_logging.html

Complete async example that sets up a metric logger, records three metrics with different reduction types (SUM, MAX, MEAN), flushes at global_step=0, and shuts down. The output shows the aggregated values after accumulating numbers 1, 2, and 3.

```python
import asyncio
from forge.observability import get_or_create_metric_logger, record_metric, Reduce

async def main():
    # Setup logger
    mlogger = await get_or_create_metric_logger(process_name="Controller")
    await mlogger.init_backends.call_one({"console": {"logging_mode": "global_reduce"}})

    # Have this in any process
    def my_fn(number):
        record_metric("my_sum_metric", number, Reduce.SUM)   # sum(1,2,3)
        record_metric("my_max_metric", number, Reduce.MAX)   # max(1,2,3)
        record_metric("my_mean_metric", number, Reduce.MEAN) # mean(1,2,3)

    # Accumulate metrics
    for number in range(1, 4): # 1, 2, 3
        my_fn(number)

    # Flush
    await mlogger.flush.call_one(global_step=0)

    # Shutdown when done
    await mlogger.shutdown.call_one()

if __name__ == "__main__":
    asyncio.run(main())

```

```text
=== [GlobalReduce] - METRICS STEP 0 ===
my_sum_metric:  6.0
my_max_metric:  3.0
my_mean_metric: 2.0

```

--------------------------------

### Run SFT Training

Source: https://meta-pytorch.org/torchforge/main/getting_started.html

Command to run supervised fine-tuning (SFT) on Llama 3 8B. Requires 2+ GPUs. Hugging Face authentication is required; run 'huggingface-cli login' first if not already done.

```bash
 python -m apps.sft.main --config apps/sft/llama3_8b.yaml
```

--------------------------------

### Get internal state with @endpoint _get_internal_state

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Returns comprehensive internal state for testing/debugging. Updates service metrics first. Not intended for production use.

```python
    @endpoint
    def _get_internal_state(self) -> dict:
        """
        Gets comprehensive internal state for testing purposes.

        This is intended for testing/debugging only, it should not
        be relied upon in actual production code.
        """
        # Ensure metrics are up to date
        self._update_service_metrics()
```

--------------------------------

### Create Generator service with .options().as_service() and route

Source: https://meta-pytorch.org/torchforge/main/tutorials/zero-to-forge/2_Forge_Internals.html

Creates a Generator service using the .options().as_service() pattern from the GRPO implementation. It configures engine and sampling parameters, then routes a prompt and prints the response, followed by cleanup.

```python
from forge.actors.generator import Generator

model = "Qwen/Qwen3-1.7B"

generator = await Generator.options(
    procs=1,
    with_gpus=True,
    num_replicas=1
).as_service(
    engine_config={
        "model": model,
        "tensor_parallel_size": 1,
        "pipeline_parallel_size": 1,
        "enforce_eager": False
    },
    sampling_config={
        "n": 1,
        "max_tokens": 16,
        "temperature": 1.0,
        "top_p": 1.0
    }
)

prompt = "What is 3 + 5?"
responses = await generator.generate.route(prompt)
print(f"Response: {responses[0].text}")

# Cleanup when done
await generator.shutdown()
```

--------------------------------

### Configure Generator.options with procs, num_replicas, with_gpus

Source: https://meta-pytorch.org/torchforge/main/tutorials/zero-to-forge/2_Forge_Internals.html

Shows the ServiceConfig pattern from apps/grpo/main.py, configuring processes, replicas, and GPU allocation. The commented lines list additional options like hosts.

```python
Generator.options(
    procs=1,           # Processes per replica
    num_replicas=4,    # Number of replicas
    with_gpus=True     # Allocate GPUs
    # Other available options:
    # hosts=None   #  the number of remote hosts used per replica
)
```

--------------------------------

### Get replica for session

Source: https://meta-pytorch.org/torchforge/main/_modules/forge/controller/service/service.html

Returns a replica for the given session ID. If no session is provided, uses the default router; otherwise uses the session router.

```python
healthy_replicas = [r for r in self._replicas if r.healthy]
if sess_id is None:
    # No session, use the default router
    return self._default_router.get_replica(healthy_replicas)

return self._session_router.get_replica(
    healthy_replicas, sess_id, self._session_replica_map
)
```

--------------------------------

### TitanTrainer

Source: https://meta-pytorch.org/torchforge/main/api_trainer.html

The TitanTrainer class provides methods and properties for configuring and running training jobs. It includes methods for setup, training steps, checkpointing, and weight management.

```APIDOC
## TitanTrainer

### Description
The TitanTrainer class is the main entry point for training with TorchForge. It provides methods for setup, training steps, checkpointing, and weight management.

### Methods
- `forward_backward()` - Performs forward and backward pass.
- `loss()` - Computes the loss.
- `train_step` - Performs a single training step.
- `setup` - Sets up the trainer.
- `load` - Loads a checkpoint.
- `save` - Saves a checkpoint.
- `cleanup` - Cleans up resources.
- `clear_gradients` - Clears gradients.
- `compile` - Compiles the model.
- `push_weights` - Pushes weights.
- `get_config` - Gets the configuration.
- `get_status` - Gets the status.

### Properties
- `activation_checkpoint` - Activation checkpointing configuration.
- `checkpoint` - Checkpoint configuration.
- `comm` - Communication configuration.
- `job` - Job configuration.
- `lr_scheduler` - Learning rate scheduler.
- `memory_estimation` - Memory estimation configuration.
- `model` - Model configuration.
- `optimizer` - Optimizer configuration.
- `parallelism` - Parallelism configuration.
- `quantize` - Quantization configuration.
- `state_dict_key` - State dict key.
- `training` - Training configuration.
```
Tokens:
Raw