nexaroa 0.0.111__py3-none-any.whl

neuroshard/core/swarm/compute.py

@@ -0,0 +1,624 @@
+"""
+Compute Engine - Decoupled GPU Worker with Soft Overflow
+
+Implements async compute loop with:
+- Priority-based activation processing
+- Interleaved 1F1B schedule (forward/backward interleaving)
+- Soft overflow handling ("Don't Stop" logic)
+- DiLoCo-style local gradient accumulation during congestion
+
+Key Directive: "If outbound.full(): Do not await. Instead: 
+accumulate_gradient_locally() and discard the activation. 
+Treat it as a DiLoCo-style local-only training step."
+
+CRITICAL: GPU must NEVER wait for network. Compute loop must always make progress.
+"""
+
+import asyncio
+import logging
+import time
+import torch
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, Optional, Any, List, TYPE_CHECKING
+
+from neuroshard.core.swarm.buffers import (
+    ActivationBuffer,
+    OutboundBuffer,
+    ActivationPacket,
+    ActivationPriority,
+)
+
+if TYPE_CHECKING:
+    from neuroshard.core.swarm.diloco import DiLoCoTrainer
+
+logger = logging.getLogger(__name__)
+
+
+class StepOutcome(Enum):
+    """Outcome of a compute step."""
+    SENT = "sent"               # Normal: activation sent to next peer
+    LOCAL_ONLY = "local_only"   # Soft overflow: accumulated locally, activation discarded
+    DROPPED = "dropped"         # Critical overflow: couldn't even accumulate
+
+
+@dataclass
+class ComputeStats:
+    """Statistics for compute engine."""
+    total_steps: int = 0
+    forward_count: int = 0
+    backward_count: int = 0
+    local_only_steps: int = 0
+    dropped_steps: int = 0
+    total_compute_time_ms: float = 0.0
+    total_queue_time_ms: float = 0.0
+    
+    @property
+    def local_only_rate(self) -> float:
+        """Fraction of steps that were local-only due to overflow."""
+        if self.total_steps == 0:
+            return 0.0
+        return self.local_only_steps / self.total_steps
+    
+    @property
+    def drop_rate(self) -> float:
+        """Fraction of steps that were completely dropped."""
+        if self.total_steps == 0:
+            return 0.0
+        return self.dropped_steps / self.total_steps
+    
+    @property
+    def avg_compute_time_ms(self) -> float:
+        """Average compute time per step."""
+        if self.total_steps == 0:
+            return 0.0
+        return self.total_compute_time_ms / self.total_steps
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "total_steps": self.total_steps,
+            "forward_count": self.forward_count,
+            "backward_count": self.backward_count,
+            "local_only_steps": self.local_only_steps,
+            "dropped_steps": self.dropped_steps,
+            "local_only_rate": self.local_only_rate,
+            "drop_rate": self.drop_rate,
+            "avg_compute_time_ms": self.avg_compute_time_ms,
+        }
+
+
+class ComputeEngine:
+    """
+    Decoupled GPU compute worker with Soft Overflow handling.
+    
+    Architecture:
+    
+        ┌─────────────┐     ┌─────────────┐     ┌─────────────┐
+        │ InboundQueue│ ──→ │ GPU Compute │ ──→ │OutboundQueue│
+        │ (Priority)  │     │ (Never Wait)│     │ (Async Send)│
+        └─────────────┘     └─────────────┘     └─────────────┘
+    
+    Key Behaviors:
+    
+    1. PULLS from inbound buffer (priority queue)
+    2. COMPUTES forward/backward pass on GPU
+    3. PUSHES to outbound buffer (if not congested)
+    
+    CRITICAL: Never waits for network - GPU must never stall!
+    
+    Soft Overflow Logic ("Don't Stop" Mechanism):
+    =============================================
+    When outbound buffer is full (network congestion):
+    
+    1. DO NOT await outbound.put() - this would stall GPU
+    2. Instead, accumulate gradients locally (DiLoCo style)
+    3. Discard activation (don't try to send)
+    4. Continue processing next packet
+    
+    Rationale: Better to treat a step as "local-only training" than to
+    halt the GPU waiting for network. DiLoCo outer optimizer syncs later.
+    
+    Interleaved 1F1B Schedule:
+    ==========================
+    For 4 micro-batches: F0 F1 F2 F3 B0 F4 B1 F5 B2 F6 B3 ...
+    
+    Start backward passes BEFORE all forwards complete, overlapping
+    backward compute with forward network latency.
+    """
+    
+    # Soft overflow thresholds
+    OUTBOUND_SOFT_LIMIT = 0.9    # Start soft overflow at 90% full
+    OUTBOUND_HARD_LIMIT = 0.99   # Hard limit - must discard
+    
+    # Scheduling parameters
+    DEFAULT_NUM_MICRO_BATCHES = 4
+    DEFAULT_WARMUP_STEPS = 4     # Forward steps before interleaving
+    
+    def __init__(
+        self,
+        model: Any,  # DynamicNeuroLLM
+        inbound: ActivationBuffer,
+        outbound: OutboundBuffer,
+        diloco_trainer: Optional['DiLoCoTrainer'] = None,
+        num_micro_batches: int = DEFAULT_NUM_MICRO_BATCHES,
+        node_id: str = "",
+    ):
+        """
+        Initialize compute engine.
+        
+        Args:
+            model: Neural network model (DynamicNeuroLLM)
+            inbound: Buffer for incoming activations
+            outbound: Buffer for outgoing activations
+            diloco_trainer: Optional DiLoCo trainer for local accumulation
+            num_micro_batches: Number of micro-batches for 1F1B schedule
+            node_id: This node's identifier
+        """
+        self.model = model
+        self.inbound = inbound
+        self.outbound = outbound
+        self.diloco = diloco_trainer
+        self.num_micro_batches = num_micro_batches
+        self.node_id = node_id
+        
+        # Device handling
+        self.device = getattr(model, 'device', 'cpu')
+        if hasattr(model, 'device'):
+            self.device = model.device
+        elif torch.cuda.is_available():
+            self.device = torch.device('cuda')
+        else:
+            self.device = torch.device('cpu')
+        
+        # Interleaved 1F1B state
+        self.pending_backwards: Dict[int, ActivationPacket] = {}
+        self.saved_activations: Dict[int, torch.Tensor] = {}  # For backward pass
+        
+        # Soft overflow state
+        self.local_gradient_buffer: Dict[str, torch.Tensor] = {}
+        
+        # Statistics
+        self.stats = ComputeStats()
+        
+        # State
+        self.running = False
+        self._task: Optional[asyncio.Task] = None
+        
+        # Callbacks
+        self._on_forward_complete: Optional[callable] = None
+        self._on_backward_complete: Optional[callable] = None
+        
+    def _check_outbound_pressure(self) -> str:
+        """
+        Check outbound buffer pressure level.
+        
+        Returns:
+            "ok" - can send normally
+            "soft_overflow" - buffer almost full, use local accumulation
+            "hard_overflow" - buffer completely full, must discard
+        """
+        fill_rate = self.outbound.fill_rate
+        
+        if fill_rate >= self.OUTBOUND_HARD_LIMIT:
+            return "hard_overflow"
+        elif fill_rate >= self.OUTBOUND_SOFT_LIMIT:
+            return "soft_overflow"
+        else:
+            return "ok"
+    
+    async def start(self):
+        """Start the compute loop."""
+        if self.running:
+            return
+            
+        self.running = True
+        self._task = asyncio.create_task(self.run())
+        logger.info(f"ComputeEngine started on {self.device}")
+        
+    async def stop(self):
+        """Stop the compute loop gracefully."""
+        self.running = False
+        
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+            self._task = None
+            
+        logger.info("ComputeEngine stopped")
+        
+    async def run(self):
+        """
+        Main compute loop with Interleaved 1F1B schedule and Soft Overflow.
+        
+        This is the heart of the async engine. It:
+        1. Pulls packets from inbound queue
+        2. Processes forward/backward with interleaving
+        3. Handles network congestion gracefully
+        """
+        logger.info("ComputeEngine run loop started")
+        
+        while self.running:
+            try:
+                # Non-blocking get from inbound buffer
+                packet = self.inbound.get(timeout=0.01)
+                
+                if packet is None:
+                    # Buffer empty - GPU potentially starved
+                    # Small sleep to avoid busy-wait
+                    await asyncio.sleep(0.001)
+                    continue
+                
+                # Process the packet
+                await self._process_packet(packet)
+                
+                # Interleaved 1F1B: After warmup, interleave backwards
+                if self.stats.forward_count >= self.num_micro_batches:
+                    await self._try_interleaved_backward()
+                    
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"ComputeEngine error: {e}", exc_info=True)
+                await asyncio.sleep(0.1)  # Prevent tight error loop
+                
+        # Cleanup
+        await self._flush_pending()
+        logger.info("ComputeEngine run loop ended")
+        
+    async def _process_packet(self, packet: ActivationPacket):
+        """Process a single activation packet."""
+        self.stats.total_steps += 1
+        start_time = time.time()
+        
+        # Track queue wait time
+        queue_time = (start_time - packet.timestamp) * 1000
+        self.stats.total_queue_time_ms += queue_time
+        
+        if packet.is_backward:
+            outcome = await self._process_backward(packet)
+        else:
+            outcome = await self._process_forward_with_overflow(packet)
+            
+        # Track compute time
+        compute_time = (time.time() - start_time) * 1000
+        self.stats.total_compute_time_ms += compute_time
+        
+        # Update stats based on outcome
+        if outcome == StepOutcome.LOCAL_ONLY:
+            self.stats.local_only_steps += 1
+        elif outcome == StepOutcome.DROPPED:
+            self.stats.dropped_steps += 1
+            
+        # Periodic logging
+        if self.stats.total_steps % 100 == 0:
+            self._log_stats()
+            
+    def _log_stats(self):
+        """Log current statistics."""
+        logger.info(
+            f"ComputeEngine: steps={self.stats.total_steps}, "
+            f"forward={self.stats.forward_count}, backward={self.stats.backward_count}, "
+            f"local_only={self.stats.local_only_rate:.1%}, "
+            f"dropped={self.stats.drop_rate:.1%}, "
+            f"avg_compute={self.stats.avg_compute_time_ms:.1f}ms"
+        )
+        
+    async def _process_forward_with_overflow(self, packet: ActivationPacket) -> StepOutcome:
+        """
+        Process forward pass with soft overflow handling.
+        
+        The "Don't Stop" Logic:
+        1. Always compute forward pass (GPU never waits)
+        2. Check outbound pressure AFTER compute
+        3. If congested: accumulate locally, skip sending
+        4. If ok: queue for outbound
+        """
+        # ALWAYS compute forward - GPU must never stall
+        try:
+            with torch.no_grad() if not packet.requires_grad else torch.enable_grad():
+                input_tensor = packet.tensor_data.to(self.device)
+                
+                # Forward through model layers
+                if hasattr(self.model, 'forward_my_layers'):
+                    output = self.model.forward_my_layers(input_tensor)
+                else:
+                    output = self.model(input_tensor)
+                    
+                # Save activation for potential backward pass
+                if packet.requires_grad:
+                    self.saved_activations[packet.micro_batch_id] = output.detach().clone()
+                    
+        except Exception as e:
+            logger.error(f"Forward pass error: {e}")
+            return StepOutcome.DROPPED
+            
+        self.stats.forward_count += 1
+        
+        # Check backpressure AFTER compute
+        pressure = self._check_outbound_pressure()
+        
+        if pressure == "ok":
+            # Normal path: queue activation for sending
+            return await self._queue_forward_output(packet, output)
+        elif pressure == "soft_overflow":
+            # SOFT OVERFLOW: Network congested
+            return self._handle_soft_overflow(packet, output)
+        else:
+            # HARD OVERFLOW: Critical congestion
+            return self._handle_hard_overflow(packet, output)
+            
+    async def _queue_forward_output(
+        self, 
+        packet: ActivationPacket, 
+        output: torch.Tensor
+    ) -> StepOutcome:
+        """Queue forward output for sending."""
+        # Determine next layer
+        if hasattr(self.model, 'my_layer_ids'):
+            next_layer = max(self.model.my_layer_ids) + 1
+        else:
+            next_layer = packet.target_layer + 1
+            
+        outbound_packet = ActivationPacket(
+            priority=packet.priority,
+            session_id=packet.session_id,
+            micro_batch_id=packet.micro_batch_id,
+            tensor_data=output.cpu(),
+            source_node=self.node_id,
+            target_layer=next_layer,
+            requires_grad=packet.requires_grad,
+        )
+        
+        # Non-blocking put with short timeout
+        try:
+            await asyncio.wait_for(
+                self.outbound.put(outbound_packet),
+                timeout=0.01  # 10ms max wait
+            )
+            return StepOutcome.SENT
+        except asyncio.TimeoutError:
+            # Couldn't send in time - fall through to soft overflow
+            return self._handle_soft_overflow(packet, output)
+        except asyncio.QueueFull:
+            return self._handle_soft_overflow(packet, output)
+            
+    def _handle_soft_overflow(
+        self, 
+        packet: ActivationPacket, 
+        output: torch.Tensor
+    ) -> StepOutcome:
+        """
+        Handle soft overflow: accumulate locally, discard activation.
+        
+        This implements DiLoCo "local training" behavior during congestion.
+        """
+        logger.debug(
+            f"Soft overflow at step {self.stats.total_steps}: "
+            f"accumulating locally (outbound: {self.outbound.fill_rate:.1%})"
+        )
+        
+        self.outbound.soft_overflow_count += 1
+        
+        # Accumulate gradient locally if training
+        if packet.requires_grad and self.model.training:
+            self._accumulate_local_gradient(output, packet)
+            
+        # Discard activation - don't try to send
+        del output
+        
+        return StepOutcome.LOCAL_ONLY
+        
+    def _handle_hard_overflow(
+        self, 
+        packet: ActivationPacket, 
+        output: torch.Tensor
+    ) -> StepOutcome:
+        """
+        Handle hard overflow: must drop step entirely.
+        """
+        logger.warning(
+            f"Hard overflow at step {self.stats.total_steps}: "
+            f"dropping step (outbound: {self.outbound.fill_rate:.1%})"
+        )
+        
+        self.outbound.hard_overflow_count += 1
+        
+        del output
+        return StepOutcome.DROPPED
+        
+    def _accumulate_local_gradient(
+        self, 
+        output: torch.Tensor, 
+        packet: ActivationPacket
+    ):
+        """
+        Accumulate gradient locally during soft overflow.
+        
+        This is the DiLoCo "local training" behavior - gradients are
+        accumulated locally and will be synced during next outer step.
+        """
+        if not self.model.training:
+            return
+            
+        # If we have upstream gradient, compute local gradient
+        if packet.grad_output is not None:
+            try:
+                output.backward(packet.grad_output.to(self.device))
+                
+                # Track in DiLoCo trainer if available
+                if self.diloco:
+                    self.diloco.inner_step_count += 1
+            except Exception as e:
+                logger.debug(f"Local gradient accumulation error: {e}")
+                
+    async def _process_backward(self, packet: ActivationPacket) -> StepOutcome:
+        """
+        Process backward pass.
+        
+        Backward passes must always be processed - they contain gradients
+        that need to be applied. However, gradient sending respects overflow.
+        """
+        try:
+            # Get saved activation for this micro-batch
+            saved_act = self.saved_activations.pop(packet.micro_batch_id, None)
+            
+            if saved_act is not None and packet.grad_output is not None:
+                grad_output = packet.grad_output.to(self.device)
+                
+                # Recompute forward and backward
+                saved_act.requires_grad_(True)
+                
+                # Backward through model
+                if hasattr(self.model, 'backward_my_layers'):
+                    grad_input = self.model.backward_my_layers(saved_act, grad_output)
+                else:
+                    # Standard backward
+                    saved_act.backward(grad_output)
+                    grad_input = saved_act.grad
+                    
+        except Exception as e:
+            logger.error(f"Backward pass error: {e}")
+            return StepOutcome.DROPPED
+            
+        self.stats.backward_count += 1
+        
+        # Backward passes also respect soft overflow for gradient sending
+        pressure = self._check_outbound_pressure()
+        
+        if pressure != "ok":
+            logger.debug(f"Backward gradient local accumulation")
+            return StepOutcome.LOCAL_ONLY
+            
+        return StepOutcome.SENT
+        
+    async def _try_interleaved_backward(self):
+        """
+        Interleaved 1F1B: Process oldest pending backward if available.
+        
+        This interleaves backward passes with forward passes to hide
+        network latency and maintain GPU utilization.
+        """
+        if not self.pending_backwards:
+            return
+            
+        # Get oldest micro-batch that needs backward
+        oldest_mb = min(self.pending_backwards.keys())
+        packet = self.pending_backwards.pop(oldest_mb)
+        
+        await self._process_backward(packet)
+        
+    async def _flush_pending(self):
+        """Process any remaining pending backward passes."""
+        while self.pending_backwards:
+            oldest_mb = min(self.pending_backwards.keys())
+            packet = self.pending_backwards.pop(oldest_mb)
+            await self._process_backward(packet)
+            
+    def get_stats(self) -> Dict[str, Any]:
+        """Get comprehensive engine statistics."""
+        return {
+            **self.stats.to_dict(),
+            "outbound_fill_rate": self.outbound.fill_rate,
+            "inbound_fill_rate": self.inbound.fill_rate,
+            "pending_backwards": len(self.pending_backwards),
+            "saved_activations": len(self.saved_activations),
+            "device": str(self.device),
+        }
+
+
+class InferenceEngine:
+    """
+    Simplified compute engine for inference-only workloads.
+    
+    No gradient handling, no backward passes, no soft overflow.
+    Just fast forward pass processing.
+    """
+    
+    def __init__(
+        self,
+        model: Any,
+        inbound: ActivationBuffer,
+        outbound: OutboundBuffer,
+        node_id: str = "",
+    ):
+        self.model = model
+        self.inbound = inbound
+        self.outbound = outbound
+        self.node_id = node_id
+        
+        # Device
+        self.device = getattr(model, 'device', torch.device('cpu'))
+        
+        # Stats
+        self.requests_processed = 0
+        self.total_latency_ms = 0.0
+        
+        self.running = False
+        self._task: Optional[asyncio.Task] = None
+        
+    async def start(self):
+        """Start inference loop."""
+        self.running = True
+        self._task = asyncio.create_task(self._run())
+        
+    async def stop(self):
+        """Stop inference loop."""
+        self.running = False
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+                
+    async def _run(self):
+        """Main inference loop."""
+        while self.running:
+            try:
+                packet = self.inbound.get(timeout=0.01)
+                if packet is None:
+                    await asyncio.sleep(0.001)
+                    continue
+                    
+                start = time.time()
+                
+                with torch.no_grad():
+                    input_tensor = packet.tensor_data.to(self.device)
+                    output = self.model(input_tensor)
+                    
+                # Queue output
+                outbound_packet = ActivationPacket(
+                    priority=packet.priority,
+                    session_id=packet.session_id,
+                    micro_batch_id=packet.micro_batch_id,
+                    tensor_data=output.cpu(),
+                    source_node=self.node_id,
+                    target_layer=packet.target_layer + 1,
+                )
+                
+                await self.outbound.put(outbound_packet)
+                
+                self.requests_processed += 1
+                self.total_latency_ms += (time.time() - start) * 1000
+                
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logger.error(f"Inference error: {e}")
+                
+    def get_stats(self) -> Dict[str, Any]:
+        """Get inference statistics."""
+        avg_latency = (
+            self.total_latency_ms / max(1, self.requests_processed)
+        )
+        return {
+            "requests_processed": self.requests_processed,
+            "avg_latency_ms": avg_latency,
+            "inbound_fill": self.inbound.fill_rate,
+            "outbound_fill": self.outbound.fill_rate,
+        }
+