DeepSpeed Worker

DeepSpeed Worker

Overview

The DeepSpeedWorker is a Ray remote actor that manages a single GPU for distributed training and inference within the Arctic Platform RL system. It wraps the DeepSpeed framework to provide:

• Distributed training with gradient accumulation and optimizer state management
• Forward-only inference for log-probability computation with ZeRO parameter sharding
• Weight synchronization via WeightSender for multi-GPU model updates
• Optional optimizations including gradient checkpointing, Liger kernels, and ZoRRO training

Initialization

Creating a Worker

Workers are created as Ray remote actors with rank, world size, and master port information:

@ray.remote
class DeepSpeedWorker:
    def __init__(self, rank: int, world_size: int, master_port: int):
        self.rank = rank
        self.world_size = world_size
        self.my_addr = socket.gethostname()
        self.master_addr = primary_ip()
        self.master_port = master_port
        self.engine = None

Initialization Method

Call initialize() with the master address and job configuration to set up distributed training:

def initialize(self, master_addr: str, job_config: dict) -> bool:

Parameters:

• master_addr: IP address of the master node for distributed coordination
• job_config: Dictionary containing:
  • model_name: HuggingFace model identifier
  • job_type: One of "training" or "log_prob" (forward-only)
  • ds_config: DeepSpeed configuration dict (optional)
  • ds_worker_config: Worker-specific config with options like:
    • attn_implementation: Attention implementation (default: "flash_attention_2")
    • enable_gradient_checkpointing: Enable gradient checkpointing (default: True)
    • use_liger: Apply Liger kernel optimizations (default: False)
    • zorro_train_enable: Enable ZoRRO training (default: False)
    • use_autocast: Enable torch autocast (default: False)
  • training_config: Optimizer and scheduler configuration
  • log_prob_config: Forward-only inference configuration
  • full_determinism: Enable full determinism (default: False)
  • seed: Random seed for determinism

Returns: True if initialization succeeds.

Example:

job_config = {
    "model_name": "mistralai/Mistral-7B-v0.1",
    "job_type": "training",
    "ds_config": {"train_micro_batch_size_per_gpu": 2},
    "ds_worker_config": {
        "attn_implementation": "flash_attention_2",
        "enable_gradient_checkpointing": True,
    },
    "training_config": {
        "optimizer": {"lr": 1e-5},
    },
}
worker.initialize.remote("192.168.1.100", job_config)

Training and Inference

Forward and Backward Pass

Perform a forward-backward pass with gradient accumulation:

def forward_backward(self, batch: dict) -> dict:

Parameters:

• batch: Dictionary containing:
  • "batch": Model input dict with input_ids, attention_mask, and other HuggingFace model inputs
  • "meta": Metadata including "worker_return_tensors" flag
  • "processing": Processing configuration

Returns: Dictionary with:

• "metrics": Combined metrics across micro-batches (loss, KL divergence, etc.)
• Other outputs from the pipeline (e.g., logits, loss)

Forward-Only Pass

Compute forward pass without gradients (used for sampling or reference model):

def forward_no_grad(self, batch: dict) -> dict:

Same parameters and return format as forward_backward().

Optimizer Step

Apply optimizer updates and return metrics:

def step(self) -> dict:

Returns: Dictionary containing:

• "metrics": Dict with "last_lr" and optionally "grad_norm"
• "batch": Empty dict

Advanced Operations

Log-Probability Computation

Compute per-token log probabilities for a batch (used for reference model evaluation):

def compute_log_probs(self, batch: dict) -> torch.Tensor:

Parameters:

• batch: Shard dict with encoded input

Returns: Tensor of shape [shard_B, S-1] on CPU, containing per-token log-probabilities.

Example:

# batch is a shard from data-parallel distribution
log_probs = worker.compute_log_probs.remote(batch)

Weight Management

Get weights (respects ZeRO-3 parameter sharding):

def get_weights(self) -> list[tuple[str, torch.Tensor]]:

Returns list of (parameter_name, parameter_tensor) tuples. Automatically materializes full parameters from ZeRO-3 shards.

Compute weight norm (for verification):

def weight_norm(self) -> dict:

Returns global L2 norm across all parameters (respects ZeRO-3 sharding). Used to verify weight synchronization.

Initialize weight sender (for distributed weight sync):

def init_weight_sender(self, group, schedule, master_addr, base_port, bucket_size) -> bool:

Parameters:

• group: NCCL process group for communication
• schedule: Sync schedule
• master_addr: Master node address
• base_port: Base port for weight sync
• bucket_size: Size of weight buckets for transmission

Checkpointing

Save and restore model state:

def save_checkpoint(self, path: str) -> bool:

Saves the DeepSpeed engine checkpoint to the specified path.

DeepSpeed Configuration

Training Configuration

For trainable engines (job_type != "log_prob"), the worker builds a DeepSpeed config with optimizer and scheduler:

def ds_training_config(self, job_config: dict, ds_config: dict, ds_worker_config: dict) -> dict:

Optimizer defaults:

• Type: AdamW
• Learning rate: 1e-5
• Betas: [0.9, 0.999]
• Epsilon: 1e-8
• Weight decay: 0.0

Scheduler options:

• "type": "cosine" — WarmupCosineLR
• "type": "constant" — No scheduler (default)
• Requires training_horizon > 0 and training_config with scheduler settings

Example:

training_config = {
    "optimizer": {
        "lr": 5e-5,
        "betas": [0.9, 0.999],
        "weight_decay": 0.01,
    },
    "lr_scheduler": {
        "type": "cosine",
        "warmup_ratio": 0.1,
        "min_lr_ratio": 0.1,
    },
    "training_horizon": 100000,
    "gradient_accumulation_steps": 4,
}

Inference Configuration

For forward-only engines (job_type == "log_prob"):

def ds_inference_config(self, log_prob_config: dict, ds_worker_config: dict) -> dict:

Keeps ZeRO parameter sharding but omits optimizer state, gradient accumulation, and optimizer offload.

Optional Features

Liger Kernel Integration

Enable fast fused operations (rope, RMS norm, SwigLU):

"ds_worker_config": {
    "use_liger": True,
    "attn_implementation": "flash_attention_2",
}

Applies _apply_liger_kernel_to_instance() with rope, RMS norm, and SwigLU optimizations.

ZoRRO Training

Enable ZoRRO training optimization:

"ds_worker_config": {
    "zorro_train_enable": True,
    "response_len": 256,
    "max_token_len": 2048,
    "rollout_n": 4,
    "temperature": 0.7,
    "logits_optimization": "none",  # or "recompute", "store"
    "use_unpad": True,
}

Uses Qwen3ModelOncePatcher to optimize the forward pass for RL training.

Gradient Checkpointing

Enable/disable gradient checkpointing (default: enabled):

"ds_worker_config": {
    "enable_gradient_checkpointing": True,
}

Environment Setup

The worker automatically configures distributed training environment variables:

RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, MASTER_PORT

It handles NCCL topology file cleanup for aws-ofi-nccl compatibility by removing stale inherited NCCL_TOPO_FILE values.

Memory and Debugging

Memory Tracking

Call see_memory_usage() for memory diagnostics (enabled in forward pass):

see_memory_usage("_forward_maybe_backward start", force=True)

Maximum Parameter Size

Get the largest parameter tensor size (in bytes):

def max_param_bytes(self) -> int:

Useful for estimating memory overhead.

Integration with Ray Server

The worker is typically managed by a Ray server that orchestrates:

• Multiple DeepSpeed workers for data parallelism
• ArcticInference ReplicaPools for inference
• Batch distribution and metric aggregation

See the server documentation for usage patterns.