# How `--colocate-memory-peak-device gpu` Optimizes Memory Allocation in Miles

> Discover how the --colocate-memory-peak-device gpu option in Miles optimizes GPU memory by offloading KV-cache and weights to CPU, significantly reducing peak consumption for larger models.

- Repository: [RadixArk/miles](https://github.com/radixark/miles)
- Tags: performance
- Published: 2026-09-06

---

**The `--colocate-memory-peak-device gpu` option forces Miles to store only the minimal tensors needed for the current forward/backward pass on GPU, while offloading the KV-cache and model weights to CPU, dramatically reducing peak GPU memory consumption.**

This memory optimization flag is implemented in the **Miles** training framework (radixark/miles), a distributed RLHF training system designed to train large language models on limited GPU resources. By colocating the memory peak on the GPU, Miles enables training of models that would otherwise exceed single-GPU memory capacity.

## What the Flag Controls

When `--colocate-memory-peak-device gpu` is passed, Miles enforces a strict memory management protocol in [`train.py`](https://github.com/radixark/miles/blob/main/train.py). The system guarantees that only actively-computed tensors reside on GPU, while all bulk data—including KV-cache and model weights—lives on CPU until explicitly needed.

This differs from the default behavior where weights and caches may persist on GPU across iterations, causing memory accumulation.

## Prerequisites and Validation

The flag does not work in isolation. According to the source code in [`miles/utils/arguments.py`](https://github.com/radixark/miles/blob/main/miles/utils/arguments.py), enabling `colocate-memory-peak-device gpu` triggers validation that ensures compatible configuration:

```python
assert args.offload_train and args.offload_rollout
assert not args.use_critic

```

- **`offload_train`** — Actor model weights reside on CPU during training
- **`offload_rollout`** — Rollout model weights reside on CPU during generation
- **No critic** — The value model path must be disabled to prevent additional GPU memory consumption

Additionally, the code raises `NotImplementedError` if **LoRA** is detected, as the current implementation only supports full-parameter training:

```python
if is_lora_enabled(args):
    raise NotImplementedError

```

## The Offload/Onload Cycle

The core optimization occurs in a carefully orchestrated sequence of tensor movements. Here is the actual control flow from [`train.py`](https://github.com/radixark/miles/blob/main/train.py):

### Before Each Rollout

```python
await inference_controller.offload_kv()        # Move KV-cache to CPU

await actor_model.onload()                     # Load minimal weights for forward pass

await inference_controller.offload_weights()   # Ensure rollout weights on CPU

```

**Why this reduces memory:** The KV-cache, which scales quadratically with sequence length, is completely evacuated from GPU. Only the specific weight shards needed for the upcoming forward computation are loaded.

### After Forward/Backward Pass

```python
await actor_model.clear_memory()               # Free all GPU buffers

await inference_controller.onload_weights()    # Prepare weights for next iteration

await offload_train()                          # Ensure training model off-GPU

```

**The "clear-then-load" pattern:** Instead of accumulating tensors across steps, GPU memory is aggressively purged. This bounds peak allocation to the size of one forward/backward pass rather than the sum of all cached states.

## Weight Update Propagation

After gradient computation, new weights must reach the rollout engines. The `update_weights` function handles this efficiently:

```python
await update_weights(actor_model, rollout_executor, rollout_id=...)

```

Weights are saved to CPU from the training process, then broadcast to rollout engines. Only actively-used engines retain the weights on GPU; others immediately offload.

## Practical Usage

### Command-Line Example

```bash
python -m miles.main.train \
    --model qwen3-6b \
    --offload-train \
    --offload-rollout \
    --colocate-memory-peak-device gpu \
    --num-rollout 100 \
    --batch-size 8

```

**Memory behavior with this configuration:**

- Model weights: Stored on CPU (~12-24 GB for 6B-7B models)
- KV-cache: Offloaded after each rollout
- GPU resident: Only current microbatch tensors (~2-4 GB)
- **Effective GPU savings:** 60-80% reduction versus non-offloaded training

### Simplified Control Flow

```python
if args.colocate_memory_peak_device == "gpu":
    # Validation: both offload flags required

    assert args.offload_train and args.offload_rollout
    
    # Training loop

    while step < max_steps:
        # Pre-rollout: minimize GPU footprint

        await inference_controller.offload_kv()
        await actor_model.onload()          # Minimal load

        await inference_controller.offload_weights()
        
        # Execute rollout and training...

        
        # Post-step: clear and reload

        await actor_model.clear_memory()    # Aggressive deallocation

        await inference_controller.onload_weights()
        await offload_train()               # Training model to CPU

```

## Key Implementation Files

| File | Role in Memory Optimization |
|------|----------------------------|
| [`train.py`](https://github.com/radixark/miles/blob/main/train.py) | Orchestrates the offload/onload sequence based on flag state |
| [`miles/utils/arguments.py`](https://github.com/radixark/miles/blob/main/miles/utils/arguments.py) | Defines and validates `--colocate-memory-peak-device gpu` |
| [`miles/ray/inference_controller.py`](https://github.com/radixark/miles/blob/main/miles/ray/inference_controller.py) | Implements `offload_kv()`, `offload_weights()`, `onload_weights()` |
| [`miles/model.py`](https://github.com/radixark/miles/blob/main/miles/model.py) | Provides `onload()`, `offload()`, `clear_memory()` methods |

## Summary

The `--colocate-memory-peak-device gpu` option optimizes GPU memory allocation in Miles by:

- **Enforcing strict offloading** — Requires `offload_train` and `offload_rollout` to move weights and KV-cache to CPU
- **Bounding peak GPU usage** — Only tensors needed for immediate computation reside on GPU
- **Implementing aggressive cleanup** — `clear_memory()` purges buffers after each step
- **Prohibiting memory-heavy configurations** — Disables critic models and LoRA adapters
- **Cycling weight locations** — Load minimal subsets for forward pass, then immediately offload

This architecture enables training 6B-7B parameter models on single 24GB GPUs that would otherwise require 40GB+ without optimization.

## Frequently Asked Questions

### What happens if I enable `--colocate-memory-peak-device gpu` without `--offload-train`?

The training process will fail with an `AssertionError`. The source code in [`train.py`](https://github.com/radixark/miles/blob/main/train.py) explicitly requires both `offload_train` and `offload_rollout` to be `True`, as the colocation strategy depends on weights being CPU-resident by default.

### Does this option work with LoRA training?

No. The implementation raises `NotImplementedError` when LoRA is detected. Full-parameter training is required because the weight update and offload logic assumes complete model tensors that can be moved between CPU and GPU, not low-rank adapter matrices.

### How much GPU memory does this actually save?

Memory savings scale with model size and sequence length. For a 7B model with 4096-token sequences, expect **60-80% GPU memory reduction** versus non-offloaded training. The KV-cache alone can consume 10-20 GB at long contexts; colocating the peak avoids persistent allocation of this buffer.

### Is there a performance penalty for the CPU-GPU transfers?

Yes, but it is typically offset by the ability to use larger batch sizes or train larger models. The overhead of `onload()`/`offload()` operations is amortized across the forward/backward computation. For throughput-critical workloads, consider using multiple GPUs with tensor parallelism instead of single-GPU colocation.