worker_extn.py

import gc
import time
import random
import sys
import numpy as np
import torch
import os
import inspect

# Make the repo root importable on the vLLM worker so `core.perturb` resolves
# regardless of the worker process cwd.
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)
from core import perturb  # fused dense-Rademacher reconstruct/switch

try:
    from vllm.forward_context import set_forward_context
except ImportError:
    set_forward_context = None

def _stateless_init_process_group(master_address, master_port, rank, world_size, device):
    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
    from vllm.distributed.utils import StatelessProcessGroup
    pg = StatelessProcessGroup.create(
        host=master_address, port=master_port, rank=rank, world_size=world_size
    )
    return PyNcclCommunicator(pg, device=device)

class WorkerExtension:
    """vLLM worker-side methods for RandOpt's fused dense-Rademacher switching.

    Fast path (speedrun):
    - configure_perturbation(noise, kernel)  -- set once at launch
    - store_base_weights()                   -- resident base copy (1x model mem)
    - perturb_self_weights(seed, sigma)      -- W = base + sigma*R(seed) (1 fused pass)
    - restore_self_weights(...)              -- NO-OP (next perturb reconstructs from base)
    - reset_to_base_weights()                -- force W = base

    Ensemble:
    - apply_perturbation(seed, sigma)        -- same fused reconstruct
    - apply_averaged_perturbations(seeds_sigmas, weights)

    Other / legacy:
    - switch_to_seed(seed_from, seed_to, sigma)  -- in-place delta switch (no-base regime)
    - update_weights_from_seeds(...)         -- legacy ES weight update (Gaussian)
    - init_inter_engine_group / broadcast_all_weights / save_self_weights_to_disk
    """
    
    def cleanup_gpu_memory(self):
        """Explicitly clean up GPU memory. Call this between iterations."""
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
        return True
    
    # Prefixes of visual encoder parameters to skip during perturbation (for VL models)
    _VISUAL_PREFIXES = ("visual.", "model.visual.")

    def _should_perturb(self, name: str) -> bool:
        """Check if a parameter should be perturbed.
        
        By default, skips visual encoder params for VL models.
        Set env PERTURB_VISUAL=1 to also perturb visual encoder.
        """
        if os.environ.get("PERTURB_VISUAL", "0") == "1":
            return True  # Perturb ALL parameters including visual encoder
        return not name.startswith(self._VISUAL_PREFIXES)

    def _set_seed(self, seed):
        # set a seed locally on the worker extension for reproducibility
        self.local_seed = seed

        # seeding
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

    # ---- perturbation config (set by core.engine.launch_engines) ----------
    def configure_perturbation(self, noise: str = "rademacher", kernel: str = None):
        """Set the noise type ('rademacher'|'gaussian') and kernel backend
        ('auto'|'triton'|'torch', or None=auto/env). Called once at launch."""
        self._noise = noise
        self._kernel = kernel
        return True

    def warmup_kernels(self):
        """Trigger Triton JIT compilation of the reconstruct/switch kernels NOW
        (at launch), so it doesn't happen mid-inference. A cold first call inside
        the sampling loop logs 'Triton JIT compilation during inference' and can
        stall under CUDA graphs; compiling against the real param dtype/device up
        front avoids that. Runs on a throwaway buffer (never touches live weights)."""
        try:
            params = list(self.model_runner.model.named_parameters())
            if not params:
                return True
            p = params[0][1]
            buf = torch.empty(1024, dtype=p.dtype, device=p.device)
            base = torch.zeros_like(buf)
            perturb.reconstruct_into(buf, base, 0, 1e-3, 0,
                                     noise=getattr(self, "_noise", "rademacher"),
                                     kernel=getattr(self, "_kernel", None))
            if getattr(self, "_noise", "rademacher") == "rademacher":
                perturb.switch_inplace(buf, 0, 1, 1e-3, 0,
                                       kernel=getattr(self, "_kernel", None))
                perturb.reconstruct_linear2_into(buf, base, 0, 5e-4, 1, 5e-4, 0,
                                                 kernel=getattr(self, "_kernel", None))
            if torch.cuda.is_available():
                torch.cuda.synchronize()
        except Exception as e:  # warmup is best-effort; never fail launch on it
            print(f"[warmup_kernels] skipped: {type(e).__name__}: {e}")
        return True

    def _reconstruct(self, seed, sigma):
        """Set every perturbable param to base + sigma*R(seed, .) in ONE fused
        pass (no noise materialised on the Triton path). One synchronize at the
        end — no per-parameter empty_cache/synchronize churn."""
        if not hasattr(self, "_base_weights"):
            self.store_base_weights()
        perturb.reconstruct_model_from_base(
            self.model_runner.model.named_parameters(),
            lambda n: self._base_weights[n],
            int(seed), float(sigma), self._should_perturb,
            noise=getattr(self, "_noise", "rademacher"),
            kernel=getattr(self, "_kernel", None),
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return True

    def perturb_self_weights(self, seed, noise_scale, negate=False):
        """Reconstruct W = base + sigma*R(seed) absolutely from the resident base
        (drift-free). Replaces the old materialise-randn-and-add. `negate` flips
        the sign of the perturbation (antithetic), realised as -sigma."""
        sigma = -float(noise_scale) if negate else float(noise_scale)
        return self._reconstruct(seed, sigma)

    def restore_self_weights(self, seed, SIGMA, negate=False):
        """No-op. With absolute reconstruction the next perturb_self_weights()
        rebuilds the live weights from the resident base, so a dedicated restore
        pass is unnecessary (this is the headline speedup: no second full pass,
        no extra RNG). Call reset_to_base_weights() to force base now."""
        return True

    def switch_to_seed(self, seed_from, seed_to, sigma):
        """In-place delta switch live += sigma*(R(seed_to)-R(seed_from)).
        Drift-prone over many hops; for the memory-constrained / no-base regime
        only. Prefer perturb_self_weights (absolute reconstruct) otherwise."""
        perturb.switch_model(
            self.model_runner.model.named_parameters(),
            int(seed_from), int(seed_to), float(sigma), self._should_perturb,
            kernel=getattr(self, "_kernel", None),
        )
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return True

    def update_weights_from_seeds(self, seeds, coeffs, alpha, population_size):
        """
        Mimics the Original implementation's update loop structure:
        Iterate Param -> Iterate Seeds -> Accumulate -> Single Update.
        """
        # seeds and coeffs should be lists of equal length
        # coeffs[i] should be: (alpha / population_size) * normalized_reward
        
        num_seeds = len(seeds)
        param_count = 0
        
        for name, p in self.model_runner.model.named_parameters():
            if not self._should_perturb(name):
                param_count += 1
                continue
            # float32
            update_accumulator = torch.zeros_like(p.data, dtype=torch.float32)
            
            for i, seed in enumerate(seeds):
                self._set_seed(seed)
                gen = torch.Generator(device=p.device)
                gen.manual_seed(int(seed))
                
                # Generate noise (in native precision, usually float16/bfloat16)
                noise = torch.randn(p.shape, dtype=p.dtype, device=p.device, generator=gen)
                
                # FIXED: Convert noise to float32 BEFORE multiplication.
                # Use in-place operation to avoid extra memory allocation
                noise_fp32 = noise.to(torch.float32)
                del noise  # Free original noise immediately
                
                # Scale in-place and accumulate
                noise_fp32.mul_(coeffs[i])
                update_accumulator.add_(noise_fp32)
                
                # Clean up immediately to avoid memory accumulation
                del noise_fp32
            
            # div by population_size multiply by alpha (scalar)
            update_accumulator.div_(population_size)
            update_accumulator.mul_(alpha)
            # Apply final update to weight (cast back to model dtype at the very end)
            p.data.add_(update_accumulator.to(p.dtype))
            
            del update_accumulator
            param_count += 1
            
            # Periodic cache clearing for large models (every 50 parameters)
            if param_count % 50 == 0:
                torch.cuda.empty_cache()
            
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        torch.cuda.empty_cache()
        gc.collect()
        return True

    def get_worker_ip(self):
        """Return the IP address of this worker's node."""
        from vllm.utils import get_ip
        return get_ip()

    def init_inter_engine_group(self, master_address: str, master_port: int, rank: int, world_size: int):
        self.inter_pg = _stateless_init_process_group(
            master_address, master_port, rank, world_size, self.device
        )
        return True

    def broadcast_all_weights(self, src_rank: int):
        for _, p in self.model_runner.model.named_parameters():
            self.inter_pg.broadcast(p, src=int(src_rank), stream=torch.cuda.current_stream())
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return True

    def save_self_weights_to_disk(self, filepath):
        state_dict_to_save = {}
        for name, p in self.model_runner.model.named_parameters():
            state_dict_to_save[name] = p.detach().cpu()
        torch.save(state_dict_to_save, filepath)
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        time.sleep(0.1)
        return True
    
    def dump_noise_for_seed(self, seed: int, out_dir: str):
        """
        Generate per-parameter noise using the same method as perturb/restore
        and save them to disk for determinism comparison.
        """
        os.makedirs(out_dir, exist_ok=True)
        noise_state = {}
        for name, p in self.model_runner.model.named_parameters():
            gen = torch.Generator(device=p.device)
            gen.manual_seed(int(seed))
            noise = torch.randn(p.shape, dtype=p.dtype, device=p.device, generator=gen)
            noise_state[name] = noise.detach().cpu()
            del noise
        torch.save(noise_state, os.path.join(out_dir, f"noise_seed_{int(seed)}.pt"))
        if torch.cuda.is_available():
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
        return True
    
    # debug
    def print_model_weights_stats(self):
        for name, p in self.model_runner.model.named_parameters():
            print(f"Param: {name}, Shape: {p.shape}")
        return True
    
    # ==================== Ensemble Methods ====================
    
    def store_base_weights(self):
        """Store a copy of current weights as base weights for ensemble."""
        self._base_weights = {}
        for name, p in self.model_runner.model.named_parameters():
            self._base_weights[name] = p.data.clone()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return True
    
    def apply_perturbation(self, seed, sigma):
        """Apply perturbation from base weights (ensemble path). Identical to
        perturb_self_weights now — a single fused reconstruct from base."""
        if not hasattr(self, '_base_weights'):
            raise RuntimeError("Must call store_base_weights first")
        return self._reconstruct(seed, sigma)
    
    def reset_to_base_weights(self):
        """Reset model weights to stored base weights."""
        if not hasattr(self, '_base_weights'):
            raise RuntimeError("Must call store_base_weights first")
        for name, p in self.model_runner.model.named_parameters():
            p.data.copy_(self._base_weights[name])
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        return True
    
    def clear_base_weights(self):
        """Free memory used by stored base weights."""
        if hasattr(self, '_base_weights'):
            del self._base_weights
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return True
    
    def apply_averaged_perturbations(self, seeds_sigmas, weights=None):
        """
        Apply the weighted average of multiple perturbations from base weights.
        This creates a single weight-averaged model from K perturbed models.
        
        Args:
            seeds_sigmas: List of (seed, sigma) tuples
            weights: Optional list of weights for each perturbation (default: equal weights)
        
        The averaged model is: W_base + sum(w_i * sigma_i * noise_i) / sum(w_i)
        """
        if not hasattr(self, '_base_weights'):
            raise RuntimeError("Must call store_base_weights first")
        
        K = len(seeds_sigmas)
        if weights is None:
            weights = [1.0 / K] * K  # Equal weights, normalized
        else:
            # Normalize weights
            total = sum(weights)
            weights = [w / total for w in weights]

        # Non-perturbed params: ensure they equal base.
        for name, p in self.model_runner.model.named_parameters():
            if not self._should_perturb(name):
                p.data.copy_(self._base_weights[name])

        # Perturbed params: W = base + sum_i (w_i * sigma_i) * R(seed_i), accumulated
        # in fp32 for precision. Uses the same global-offset Rademacher stream as
        # the reconstruct/switch kernels (so seeds match the sampling phase).
        plan = perturb.iter_perturb_params(
            self.model_runner.model.named_parameters(), self._should_perturb)
        for name, p, off in plan:
            base = self._base_weights[name]
            acc = torch.zeros_like(p.data, dtype=torch.float32)
            for (seed, sigma), weight in zip(seeds_sigmas, weights):
                signs = perturb.rademacher_signs(
                    int(seed), off, p.numel(), p.device, torch.float32).reshape(p.shape)
                acc.add_(signs, alpha=float(weight) * float(sigma))
            p.data.copy_((base.to(torch.float32) + acc).to(p.dtype))
            del acc

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        gc.collect()
        return True

    def apply_linear_combined_perturbations(self, seeds_sigmas, coeffs):
        """
        Apply an unnormalized linear combination of perturbation deltas.

        The live model becomes:

            W = W_base + sum_i coeff_i * sigma_i * R(seed_i)

        This is intentionally different from apply_averaged_perturbations:
        callers can test raw sums, norm-preserving sums, differences, and other
        weight-space merge operations without implicit coefficient normalization.
        """
        if not hasattr(self, '_base_weights'):
            raise RuntimeError("Must call store_base_weights first")
        if len(seeds_sigmas) != len(coeffs):
            raise ValueError("seeds_sigmas and coeffs must have the same length")
        if getattr(self, "_noise", "rademacher") != "rademacher":
            raise ValueError("linear perturbation merges currently support rademacher noise only")

        if len(seeds_sigmas) == 2:
            for name, p in self.model_runner.model.named_parameters():
                if not self._should_perturb(name):
                    p.data.copy_(self._base_weights[name])
            (seed1, sigma1), (seed2, sigma2) = seeds_sigmas
            scale1 = float(coeffs[0]) * float(sigma1)
            scale2 = float(coeffs[1]) * float(sigma2)
            perturb.reconstruct_model_linear2_from_base(
                self.model_runner.model.named_parameters(),
                lambda n: self._base_weights[n],
                int(seed1), scale1, int(seed2), scale2, self._should_perturb,
                kernel=getattr(self, "_kernel", None),
            )
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            return True

        for name, p in self.model_runner.model.named_parameters():
            if not self._should_perturb(name):
                p.data.copy_(self._base_weights[name])

        plan = perturb.iter_perturb_params(
            self.model_runner.model.named_parameters(), self._should_perturb)
        for name, p, off in plan:
            base = self._base_weights[name]
            acc = torch.zeros_like(p.data, dtype=torch.float32)
            for (seed, sigma), coeff in zip(seeds_sigmas, coeffs):
                signs = perturb.rademacher_signs(
                    int(seed), off, p.numel(), p.device, torch.float32).reshape(p.shape)
                acc.add_(signs, alpha=float(coeff) * float(sigma))
            p.data.copy_((base.to(torch.float32) + acc).to(p.dtype))
            del acc

        if torch.cuda.is_available():
            torch.cuda.synchronize()
        gc.collect()
        return True
    
    def get_logits_for_prompt(self, input_ids_list):
        """
        Get logits for the last token position for a batch of prompts.
        Returns logits as CPU tensors for ensemble averaging.
        
        Args:
            input_ids_list: List of input_ids (each is a list of token ids)
        
        Returns:
            List of logits tensors (vocab_size,) for each prompt
        """
        model = self.model_runner.model
        model.eval()
        
        results = []
        with torch.no_grad():
            for input_ids in input_ids_list:
                seq_len = len(input_ids)
                # vLLM V1 expects flattened tensors (not batched)
                # input_ids: (seq_len,), positions: (seq_len,)
                ids_tensor = torch.tensor(input_ids, dtype=torch.long, device=self.device)
                positions = torch.arange(seq_len, dtype=torch.long, device=self.device)
                
                # Forward pass - get logits
                # vLLM v0.11+ requires forward context
                if set_forward_context is not None and hasattr(self.model_runner, "vllm_config"):
                    with set_forward_context(attn_metadata=None, 
                                           vllm_config=self.model_runner.vllm_config):
                        outputs = model(input_ids=ids_tensor, positions=positions)
                else:
                    # Fallback for older vLLM versions
                    if 'positions' in inspect.signature(model.forward).parameters:
                        outputs = model(input_ids=ids_tensor, positions=positions)
                    else:
                        outputs = model(input_ids=ids_tensor.unsqueeze(0))
                
                # Get logits for the last position
                # outputs may have .logits attribute or be the logits tensor directly
                logits = outputs.logits if hasattr(outputs, 'logits') else outputs
                # vLLM V1: logits shape is (seq_len, vocab_size) for flattened input
                # or (batch, seq_len, vocab_size) for batched input
                if logits.ndim == 2:
                    # Flattened: (seq_len, vocab_size)
                    last_logits = logits[-1, :].cpu()
                else:
                    # Batched: (batch, seq_len, vocab_size)
                    last_logits = logits[0, -1, :].cpu()
                results.append(last_logits)
                
                del ids_tensor, positions, outputs, logits
        
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        torch.cuda.empty_cache()
        
        return results
    
    def generate_with_logits_callback(self, input_ids, max_new_tokens, temperature=1.0):
        """
        Generate tokens step by step and return the logits at each step.
        This is for debugging/analysis - actual ensemble should use get_logits_for_prompt.
        
        Returns: (generated_ids, list_of_logits_at_each_step)
        """
        model = self.model_runner.model
        model.eval()
        
        current_ids = torch.tensor([input_ids], dtype=torch.long, device=self.device)
        all_logits = []
        
        with torch.no_grad():
            for _ in range(max_new_tokens):
                outputs = model(input_ids=current_ids)
                last_logits = outputs.logits[0, -1, :]
                all_logits.append(last_logits.cpu())
                
                # Sample next token
                if temperature > 0:
                    probs = torch.softmax(last_logits / temperature, dim=-1)
                    next_token = torch.multinomial(probs, num_samples=1)
                else:
                    next_token = last_logits.argmax(dim=-1, keepdim=True)
                
                current_ids = torch.cat([current_ids, next_token.unsqueeze(0)], dim=-1)
                
                del outputs
        
        generated = current_ids[0].cpu().tolist()
        
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        torch.cuda.empty_cache()
        
        return generated, all_logits