"""
Dynamic Model Architecture - True Decentralization

This module implements a model that grows and shrinks with the network:
- NO fixed phases or model sizes
- Model size = what the network can collectively hold
- Each node contributes based on its available memory
- More memory = more layers = more NEURO rewards

Key Concepts:
1. LAYER POOL: The network maintains a pool of layers
2. DYNAMIC ASSIGNMENT: Nodes claim layers based on their capacity
3. ORGANIC GROWTH: As more nodes join, model can have more layers
4. GRACEFUL DEGRADATION: If nodes leave, layers are redistributed

Example:
  Day 1: 10 nodes with 4GB each = 40GB total = ~10B params possible
  Day 30: 100 nodes with avg 8GB = 800GB total = ~200B params possible
  
  The model AUTOMATICALLY grows as capacity grows.
  No voting, no phases, no central coordination.
"""

import torch
import torch.nn as nn
import threading
import time
import logging
import hashlib
import math
import psutil  # For adaptive memory management
from typing import Optional, Dict, List, Tuple, Any, Set
from dataclasses import dataclass, field
from collections import defaultdict
from urllib.parse import urlparse

logger = logging.getLogger(__name__)


# Dynamic architecture - NO MORE FIXED DIMENSIONS!
# Architecture is now calculated based on network capacity

# Dynamic vocabulary - starts at 32K, expands WITHOUT LIMIT as network grows
# The embedding/lm_head grow in chunks when tokenizer vocabulary exceeds current capacity
INITIAL_VOCAB_SIZE = 32000   # Starting size (efficient for small networks)
VOCAB_GROWTH_CHUNK = 32000   # Expand by 32K at a time (efficient GPU memory alignment)

# NO HARD LIMIT - vocabulary grows with the network
# The only real constraints are:
#   - Memory: ~4KB per token (at hidden_dim=1024) or ~16KB (at hidden_dim=4096)
#   - Practical: Most use cases covered under 1M tokens
# For reference: GPT-4 ~100K, Claude ~100K, Gemini ~256K
# NeuroShard can grow FAR beyond these as a truly decentralized, ever-growing LLM
MAX_VOCAB_SIZE = None  # None = unlimited (constrained only by available memory)

# Import the new architecture scaler
from neuroshard.core.model.scaler import (
    ModelArchitecture,
    calculate_optimal_architecture,
    should_upgrade_architecture,
    estimate_memory_per_layer,
    calculate_layer_assignment,
)


@dataclass
class LayerAssignment:
    """Assignment of a layer to a node."""
    layer_id: int
    node_id: str
    node_url: str
    grpc_addr: str
    assigned_at: float = field(default_factory=time.time)
    last_heartbeat: float = field(default_factory=time.time)
    version: int = 0  # Training version


@dataclass
class NetworkCapacity:
    """Current network capacity."""
    total_nodes: int
    total_memory_mb: float
    max_layers: int  # How many layers the network can support
    assigned_layers: int  # How many layers are currently assigned
    layer_coverage: Dict[int, int]  # layer_id -> replica count


class DynamicLayerPool:
    """
    Manages the dynamic pool of model layers across the network.
    
    This is the core of true decentralization:
    - Layers are assigned based on node capacity
    - Model grows BOTH deeper AND wider as network expands
    - Architecture auto-optimizes based on scaling laws
    - No fixed model size or phases
    
    SCALABILITY CONSIDERATIONS:
    ==========================
    Small network (1-10 nodes):
    - Each node may hold ALL layers (solo training mode)
    - No layer replication needed
    - Fast startup, immediate training
    
    Medium network (10-100 nodes):
    - Layers distributed across multiple nodes
    - MIN_REPLICAS ensures redundancy
    - Pipeline inference works across nodes
    
    Large network (100-1000+ nodes):
    - Strong layer distribution
    - MAX_LAYERS_PER_NODE caps per-node load
    - Architecture can scale to 100B+ params
    """
    
    # Minimum replicas per layer for redundancy
    MIN_REPLICAS = 2
    
    # Layer heartbeat timeout
    HEARTBEAT_TIMEOUT = 120  # seconds
    
    # Maximum layers any single node can hold (prevents memory issues in large networks)
    # In small networks (< 100 nodes), this is effectively unlimited
    # In large networks, it ensures load is distributed
    MAX_LAYERS_PER_NODE = 200
    
    # Architecture recalculation triggers
    # NOTE: RECALC_INTERVAL_NODES is now DYNAMIC - see _get_recalc_interval()
    MIN_UPGRADE_IMPROVEMENT = 1.3  # Only upgrade if 30%+ better
    
    @staticmethod
    def _get_recalc_interval(node_count: int) -> int:
        """
        Get dynamic architecture recalculation interval based on network size.
        
        At small networks, recalculate more often (every node matters).
        At large networks, recalculate less often (stability).
        
        Formula: min(max(5, node_count // 10), 100)
        - 1-50 nodes: every 5 nodes
        - 51-100 nodes: every 5-10 nodes
        - 100-1000 nodes: every 10-100 nodes
        - 1000+ nodes: every 100 nodes
        """
        return min(max(5, node_count // 10), 100)
    
    def __init__(self, dht_protocol=None):
        self.dht = dht_protocol
        
        # Layer assignments
        self.layer_assignments: Dict[int, List[LayerAssignment]] = defaultdict(list)
        
        # Node capacities
        self.node_capacities: Dict[str, float] = {}  # node_id -> available_mb
        
        # DYNAMIC VOCAB: Track current vocabulary capacity for memory calculation
        # This is updated when vocab expands and affects layer assignment
        self.vocab_capacity: int = INITIAL_VOCAB_SIZE
        
        # DYNAMIC ARCHITECTURE (auto-updates as network grows)
        self.current_architecture: Optional[ModelArchitecture] = None
        self.architecture_version: int = 0
        self.last_node_count: int = 0
        
        # Legacy fields (for compatibility)
        self.current_num_layers = 0
        self.embedding_holder: Optional[str] = None
        self.lm_head_holder: Optional[str] = None
        
        # Threading
        self.lock = threading.Lock()
        
        logger.info("DynamicLayerPool initialized with dynamic width + depth scaling")
    
    def _auto_recalculate_architecture(self):
        """
        AUTOMATED architecture optimization - no human intervention needed.
        
        Calculates optimal architecture based on current network capacity
        and triggers upgrade if improvement is significant.
        """
        total_memory = sum(self.node_capacities.values())
        optimal = calculate_optimal_architecture(total_memory)
        
        if self.current_architecture is None:
            # First initialization
            self.current_architecture = optimal
            self.current_num_layers = optimal.num_layers
            self.architecture_version = 1
            logger.info(f"🚀 Initial architecture: {optimal.num_layers}L × {optimal.hidden_dim}H "
                       f"({optimal.estimate_params()/1e6:.0f}M params)")
            return
        
        # Check if upgrade is worthwhile
        should_upgrade, reason = should_upgrade_architecture(
            self.current_architecture,
            optimal,
            self.MIN_UPGRADE_IMPROVEMENT
        )
        
        if should_upgrade:
            logger.warning(f"🔄 ARCHITECTURE UPGRADE TRIGGERED!")
            logger.warning(f"   {reason}")
            logger.warning(f"   Old: {self.current_architecture.num_layers}L × {self.current_architecture.hidden_dim}H")
            logger.warning(f"   New: {optimal.num_layers}L × {optimal.hidden_dim}H")
            logger.warning(f"   Nodes will gradually migrate to new architecture on restart")
            
            # Update architecture (new nodes will use new arch)
            self.current_architecture = optimal
            self.current_num_layers = optimal.num_layers
            self.architecture_version += 1
            
            # TODO: Trigger distillation-based migration for existing nodes
            # For now, existing nodes keep their architecture until restart
        else:
            logger.debug(f"Architecture recalculation: no upgrade needed ({reason})")
    
    def register_node(
        self, 
        node_id: str, 
        node_url: str,
        grpc_addr: str,
        available_memory_mb: float,
        staked_amount: float = 0.0
    ) -> List[int]:
        """
        Register a node and assign layers based on its capacity AND stake.
        
        AUTOMATIC ARCHITECTURE SCALING:
        - Periodically recalculates optimal architecture as network grows
        - Triggers upgrades when capacity increases significantly
        - New nodes automatically use latest architecture
        
        Validator role requires:
        1. Sufficient memory (>2GB)
        2. Minimum stake (100 NEURO) - prevents Sybil attacks
        
        Returns list of layer IDs assigned to this node.
        """
        # Import validator requirements from centralized economics (with dynamic scaling!)
        from neuroshard.core.economics.constants import (
            VALIDATOR_MIN_MEMORY_MB, 
            get_dynamic_validator_stake
        )
        
        with self.lock:
            self.node_capacities[node_id] = available_memory_mb
            
            # AUTO-TRIGGER: Recalculate architecture if network grew significantly
            node_count = len(self.node_capacities)
            recalc_interval = self._get_recalc_interval(node_count)
            if (node_count - self.last_node_count) >= recalc_interval:
                self._auto_recalculate_architecture()
                self.last_node_count = node_count
            
            # Ensure we have an architecture
            if self.current_architecture is None:
                self._auto_recalculate_architecture()
            
            # Calculate how many layers this node can hold
            # Uses current architecture's dimensions (dynamic!)
            # DEVICE-AWARE safety factors
            # With gradient checkpointing always enabled, CPU can use higher factor
            device_type = getattr(self, '_device_hint', 'cpu')
            if device_type == 'cuda':
                safety_factor = 0.6  # GPU: efficient memory usage
            elif device_type == 'mps':
                safety_factor = 0.5  # Apple Silicon: moderate overhead
            else:
                safety_factor = 0.5  # CPU: increased from 0.3 (checkpointing reduces overhead)
            
            # SMART DISTRIBUTED TRAINING DESIGN:
            # Goal: Enable meaningful training for ANY network size
            #
            # KEY INSIGHT: When a "full node" exists (has embedding + LM head), 
            # new nodes should NOT become Drivers (would create broken overlap).
            # Instead, they should become WORKER+VALIDATOR to enable proper pipeline.
            #
            # Network States:
            # 1. Empty network → First node becomes DRIVER, grows to full node
            # 2. One full node exists → New node becomes WORKER+VALIDATOR (pipeline!)
            # 3. Multiple partial nodes → Fill based on MIN_REPLICAS
            
            # FULLY DECENTRALIZED: Discover network state from DHT ONLY (no tracker fallback!)
            # Each node has its own LOCAL layer_pool, so we need DHT for network-wide view.
            # BUT: DHT may have STALE data from previous runs - don't blindly trust it!
            dht_layers = set()
            
            # DHT discovery (P2P must be connected BEFORE start() for this to work!)
            if self.dht:
                dht_layers = self._discover_network_layers_from_dht()
                if dht_layers:
                    highest_layer = max(dht_layers)
                    # SANITY CHECK: Only expand if DHT layers are "reasonable"
                    max_reasonable = max(32, self.current_num_layers * 2)
                    if highest_layer >= self.current_num_layers and highest_layer < max_reasonable:
                        self.current_num_layers = highest_layer + 1
                        logger.info(f"DHT discovery: network has {self.current_num_layers} layers")
                    elif highest_layer >= max_reasonable:
                        logger.warning(f"DHT shows {highest_layer + 1} layers but seems stale")
                        logger.warning(f"Ignoring stale DHT data - will use checkpoint layer count")
                else:
                    logger.info("DHT: No existing layers found - this may be first node or peers not yet discovered")
            else:
                logger.warning("DHT not available - layer assignment will be solo mode")
            
            driver_count = len(self.layer_assignments.get(0, []))
            validator_layer = max(0, self.current_num_layers - 1) if self.current_num_layers > 0 else 0
            validator_count = len(self.layer_assignments.get(validator_layer, []))
            
            # CRITICAL: Check if a FULL NODE exists (FULLY DECENTRALIZED - DHT only!)
            # A full node has BOTH layer 0 (embedding) AND last layer (LM head)
            has_full_node = False
            full_node_id = None
            
            # Check local assignments first
            if self.current_num_layers > 0:
                layer_0_holders = {a.node_id for a in self.layer_assignments.get(0, [])}
                last_layer_holders = {a.node_id for a in self.layer_assignments.get(validator_layer, [])}
                full_node_ids = layer_0_holders & last_layer_holders  # Intersection
                if full_node_ids:
                    has_full_node = True
                    full_node_id = next(iter(full_node_ids))
                    logger.info(f"Full node detected (local): {full_node_id[:8]}... (has layers 0-{validator_layer})")
            
            # Also check DHT - if both layer 0 AND last layer exist, a full node likely exists
            if not has_full_node and dht_layers:
                if 0 in dht_layers and validator_layer in dht_layers:
                    has_full_node = True
                    full_node_id = "unknown_from_dht"
                    logger.info(f"Full node detected (DHT): network has layers 0-{validator_layer}")
            
            # ROLE ASSIGNMENT PRIORITY (redesigned for proper distributed training):
            # 
            # If NO full node exists:
            #   1. First node → DRIVER (will grow to full node)
            #   2. Additional nodes → fill based on MIN_REPLICAS
            #
            # If a FULL NODE exists:
            #   - DON'T create another Driver (would overlap and break training!)
            #   - New nodes become WORKER+VALIDATOR for proper pipeline
            #   - This creates: FullNode[0-N] → NewNode[N+1 to Last] pipeline
            
            if has_full_node:
                # A full node exists - new nodes should NOT overlap with it
                # Become WORKER+VALIDATOR to enable pipeline training!
                if node_id == full_node_id:
                    # This IS the full node re-registering
                    role_hint = "DRIVER"  # Keep as driver (already full)
                    needs_embedding = True
                    logger.info(f"This is the full node - keeping as DRIVER")
                else:
                    # New node joining a network with a full node
                    # Become WORKER+VALIDATOR: get last layer + work backwards
                    # This enables pipeline: FullNode[embedding→layers] → Us[layers→LM_head]
                    role_hint = "WORKER+VALIDATOR"
                    needs_embedding = False  # Don't need embedding, saves memory!
                    logger.info(f"Full node exists - becoming WORKER+VALIDATOR for pipeline training")
            else:
                # No full node - use standard role assignment
                needs_more_drivers = driver_count < self.MIN_REPLICAS
                needs_more_validators = validator_count < self.MIN_REPLICAS
                
                if needs_more_drivers:
                    # Network needs more Drivers for MIN_REPLICAS
                    role_hint = "DRIVER"
                    needs_embedding = True
                    logger.info(f"Network needs more Drivers ({driver_count}/{self.MIN_REPLICAS})")
                elif needs_more_validators and self.current_num_layers > 1:
                    # Network needs more Validators for MIN_REPLICAS
                    role_hint = "WORKER+VALIDATOR"
                    needs_embedding = False
                    logger.info(f"Network needs more Validators ({validator_count}/{self.MIN_REPLICAS})")
                else:
                    # Network has enough of both → become Worker for layer redundancy
                    role_hint = "WORKER"
                    needs_embedding = False
                    logger.info(f"Network has enough Drivers ({driver_count}) and Validators ({validator_count})")
            
            max_layers_for_node = calculate_layer_assignment(
                available_memory_mb,
                self.current_architecture,
                safety_factor=safety_factor,
                vocab_capacity=self.vocab_capacity,
                training_mode=True,
                needs_embedding=needs_embedding
            )
            
            logger.info(f"Layer calculation ({role_hint}): {available_memory_mb:.0f}MB × {safety_factor} safety = {max_layers_for_node} layers "
                       f"(needs_embedding={needs_embedding}, drivers={driver_count}, validators={validator_count})")
            
            # SCALABILITY: Apply MAX_LAYERS_PER_NODE cap in large networks
            # This prevents single nodes from hogging all layers and ensures
            # load distribution as the network grows
            node_count = len(self.node_capacities)
            if node_count > 100:
                # In large networks, cap layers per node
                max_layers_for_node = min(max_layers_for_node, self.MAX_LAYERS_PER_NODE)
                logger.debug(f"Large network ({node_count} nodes): capped to {max_layers_for_node} layers")
            
            if max_layers_for_node < 1:
                logger.warning(f"Node {node_id[:8]}... has insufficient memory for even 1 layer")
                return []
            
            # Find layers that need more replicas
            assigned_layers = []
            
            # SCALABILITY STRATEGY:
            # High-capacity nodes (>8GB) are prioritized for Layer 0 (Driver) and Last Layer (Validator)
            # This creates parallel pipelines ("Training Gangs")
            is_high_capacity = available_memory_mb > 8000
            is_medium_capacity = available_memory_mb > VALIDATOR_MIN_MEMORY_MB
            
            # Count current validators for DYNAMIC stake requirement
            num_drivers = len(self.layer_assignments[0])
            num_validators = len(self.layer_assignments[max(0, self.current_num_layers - 1)])
            
            # VALIDATOR ELIGIBILITY: Dynamic stake based on network size!
            # - Few validators (1-10): 100 NEURO (accessible for bootstrap)
            # - Many validators (1000+): 2500 NEURO (security at scale)
            required_validator_stake = get_dynamic_validator_stake(num_validators)
            has_validator_stake = staked_amount >= required_validator_stake
            
            # DISTRIBUTED TRAINING ASSIGNMENT:
            # Based on role_hint from capacity calculation, assign appropriate layers
            
            last_layer = max(0, self.current_num_layers - 1) if self.current_num_layers > 0 else 0
            
            # Determine role based on network needs and node capacity
            if role_hint == "DRIVER":
                # Become Driver: get Layer 0 + as many layers as we can
                should_be_driver = True
                should_be_validator = False  # Driver doesn't need to also be Validator
                logger.info(f"Node {node_id[:8]}... assigned as DRIVER (first node, will embed data)")
                
            elif role_hint == "WORKER+VALIDATOR":
                # Become Worker+Validator: skip Layer 0, get middle + last layer
                # This is the KEY for distributed training - enables pipeline with loss computation
                should_be_driver = False
                should_be_validator = True  # Need LM head to compute loss!
                # For bootstrap, allow Validator without stake requirement
                has_validator_stake = True  # Bootstrap mode - stake checked later
                logger.info(f"Node {node_id[:8]}... assigned as WORKER+VALIDATOR (will compute loss, enable distributed training)")
                
            else:  # "WORKER"
                # Become Worker: middle layers only (redundancy)
                should_be_driver = False
                should_be_validator = False
                logger.info(f"Node {node_id[:8]}... assigned as WORKER (middle layers, redundancy)")
            
            # Assign Layer 0 (Driver) - if we should be Driver
            if should_be_driver and len(assigned_layers) < max_layers_for_node:
                if not any(a.node_id == node_id for a in self.layer_assignments[0]):
                    self._assign_layer(0, node_id, node_url, grpc_addr)
                    assigned_layers.append(0)
                    # Ensure current_num_layers accounts for layer 0
                    if self.current_num_layers == 0:
                        self.current_num_layers = 1
            
            # 2. DISTRIBUTED LAYER ASSIGNMENT
            # Based on role, fill appropriate layers for pipeline parallelism
            is_driver = 0 in assigned_layers
            last_layer = max(0, self.current_num_layers - 1)
            
            if role_hint == "WORKER+VALIDATOR":
                # CRITICAL: Assign LAST layer FIRST (for LM head) to ensure we can compute loss
                if len(assigned_layers) < max_layers_for_node:
                    if last_layer not in assigned_layers and not any(a.node_id == node_id for a in self.layer_assignments[last_layer]):
                        self._assign_layer(last_layer, node_id, node_url, grpc_addr)
                        assigned_layers.append(last_layer)
                        logger.info(f"Node {node_id[:8]}... assigned last layer {last_layer} (Validator role)")
                
                # Then fill middle layers from the END (closer to Validator)
                # This creates contiguous layer ranges: Driver has 0-N, Validator has M-31
                for layer_id in range(last_layer - 1, 0, -1):  # Reverse order, skip layer 0
                    if len(assigned_layers) >= max_layers_for_node:
                        break
                    if layer_id not in assigned_layers and not any(a.node_id == node_id for a in self.layer_assignments[layer_id]):
                        self._assign_layer(layer_id, node_id, node_url, grpc_addr)
                        assigned_layers.append(layer_id)
            else:
                # Driver or Worker: fill from layer 1 onwards
                layers_to_check = range(1, self.current_num_layers)
                    
                for layer_id in layers_to_check:
                    if len(assigned_layers) >= max_layers_for_node:
                        break
                        
                    current_replicas = len(self.layer_assignments[layer_id])
                    if current_replicas < self.MIN_REPLICAS:
                        if layer_id not in assigned_layers and not any(a.node_id == node_id for a in self.layer_assignments[layer_id]):
                            self._assign_layer(layer_id, node_id, node_url, grpc_addr)
                            assigned_layers.append(layer_id)
                
                # 3. Assign Last Layer (Validator) if Driver with extra capacity
                if should_be_validator and len(assigned_layers) < max_layers_for_node:
                    if last_layer not in assigned_layers and not any(a.node_id == node_id for a in self.layer_assignments[last_layer]):
                        self._assign_layer(last_layer, node_id, node_url, grpc_addr)
                        assigned_layers.append(last_layer)

            # 4. Calculate remaining capacity
            remaining_capacity = max_layers_for_node - len(assigned_layers)
            
            # 5. DHT layer discovery already done above (at role assignment)
            # dht_layers variable is already populated
            
            # 6. If we still have capacity, grow the model
            if remaining_capacity > 0:
                # CAP MODEL GROWTH for solo/early network as safety net
                # Even with gradient checkpointing, too many layers cause OOM due to:
                # - Optimizer states (2x model size for Adam)
                # - Gradient accumulation during backward pass
                # - PyTorch memory fragmentation
                # 32 layers is safe for most devices (~200M params with 512 hidden)
                MAX_SOLO_LAYERS = 32  # Conservative cap for solo node training
                total_nodes = len(set(
                    a.node_id for assignments in self.layer_assignments.values() 
                    for a in assignments
                ))
                
                if total_nodes <= 2:  # Solo or near-solo mode
                    max_growth = max(0, MAX_SOLO_LAYERS - len(assigned_layers))
                    if remaining_capacity > max_growth:
                        logger.warning(f"[SOLO MODE] Capping growth from {remaining_capacity} to {max_growth} layers "
                                      f"(MAX_SOLO_LAYERS={MAX_SOLO_LAYERS} prevents OOM)")
                        remaining_capacity = max_growth
                
                # Add new layers to grow the model
                new_layers = self._grow_model(remaining_capacity, node_id, node_url, grpc_addr)
                assigned_layers.extend(new_layers)
            
            # Handle embedding and LM head tracking
            # Any node with Layer 0 has embedding
            if 0 in assigned_layers:
                # Update tracking (just keeps one for reference, but multiple exist)
                self.embedding_holder = node_id 
                logger.info(f"Node {node_id[:8]}... became a Driver (Layer 0)")
            
            # Any node with highest assigned layer becomes Validator (LM head holder)
            # CRITICAL FIX: Use max(assigned_layers), NOT current_num_layers - 1
            # This handles the case where checkpoint has fewer layers than stale DHT data
            if assigned_layers:
                actual_last_layer = max(assigned_layers)
                # Check if this node holds the highest layer in the network
                # OR if there's no other holder for this layer yet
                if not self.lm_head_holder or actual_last_layer >= (self.current_num_layers - 1):
                    self.lm_head_holder = node_id
                    self.current_num_layers = actual_last_layer + 1  # Update to match reality
                    logger.info(f"Node {node_id[:8]}... became a Validator (Layer {actual_last_layer})")
            
            # EARLY NETWORK NOTICE: When there are <10 nodes, each must hold many/all layers
            # This is TEMPORARY - as more nodes join, layers will be distributed
            if len(assigned_layers) > 50:
                logger.warning(f"Node {node_id[:8]}... holding {len(assigned_layers)} layers due to low network size")
                logger.warning(f"This is temporary - as more nodes join, model will shard across network")
            
            logger.info(f"Node {node_id[:8]}... registered: {len(assigned_layers)} layers assigned "
                       f"(capacity: {max_layers_for_node} layers, {available_memory_mb:.0f}MB)")
            
            return assigned_layers
    
    def _assign_layer(self, layer_id: int, node_id: str, node_url: str, grpc_addr: str):
        """Assign a layer to a node."""
        assignment = LayerAssignment(
            layer_id=layer_id,
            node_id=node_id,
            node_url=node_url,
            grpc_addr=grpc_addr,
        )
        self.layer_assignments[layer_id].append(assignment)
        
        # Announce to DHT
        if self.dht:
            try:
                import json
                key = f"layer_{layer_id}"
                current = self.dht.lookup_value(key)
                holders = json.loads(current) if current else []
                if grpc_addr not in holders:
                    holders.append(grpc_addr)
                self.dht.store(key, json.dumps(holders))
            except Exception as e:
                logger.debug(f"DHT announce failed: {e}")
    
    def _discover_network_layers_from_dht(self) -> Set[int]:
        """
        Query DHT to discover which layers exist in the network.
        
        DECENTRALIZED COORDINATION:
        - Each node announces "layer_X" to DHT when it holds layer X
        - New nodes query DHT to see what layers already exist
        - This prevents layer overlap without centralized coordination
        """
        discovered_layers = set()
        
        if not self.dht:
            return discovered_layers
        
        try:
            import hashlib
            # Query for layers 0-1000 (reasonable max)
            # DHT lookup is fast - O(log N) hops
            for layer_id in range(min(1000, self.current_num_layers + 100)):
                key_string = f"layer_{layer_id}"
                # Hash the key string to get integer key (same as announce())
                key_int = int(hashlib.sha1(key_string.encode()).hexdigest(), 16)
                try:
                    value = self.dht.lookup_value(key_int)
                    if value:
                        discovered_layers.add(layer_id)
                except Exception:
                    continue
            
            if discovered_layers:
                logger.info(f"DHT layer discovery: found {len(discovered_layers)} layers "
                           f"(range: {min(discovered_layers)}-{max(discovered_layers)})")
        except Exception as e:
            logger.debug(f"DHT layer discovery failed: {e}")
        
        return discovered_layers
    
    def _grow_model(
        self,
        num_new_layers: int,
        node_id: str,
        node_url: str,
        grpc_addr: str
    ) -> List[int]:
        """
        Grow the model by adding new layers.

        This is how the model organically grows with the network!
        """
        new_layers = []
        
        for _ in range(num_new_layers):
            new_layer_id = self.current_num_layers
            self._assign_layer(new_layer_id, node_id, node_url, grpc_addr)
            new_layers.append(new_layer_id)
            self.current_num_layers += 1
        
        if new_layers:
            logger.info(f"Model grew: now {self.current_num_layers} layers "
                       f"(added layers {new_layers[0]}-{new_layers[-1]})")
        
        return new_layers
    
    def upgrade_to_validator(self, node_id: str, node_url: str, grpc_addr: str) -> bool:
        """
        Upgrade a node to Validator role (assign LM head) when stake requirement is met.
        
        This is called when a node stakes enough NEURO to become a Validator.
        No restart required - the node's role is upgraded dynamically.
        
        Returns True if upgrade was successful.
        """
        from neuroshard.core.economics.constants import VALIDATOR_MIN_MEMORY_MB
        
        with self.lock:
            # Check if node has sufficient memory
            memory = self.node_capacities.get(node_id, 0)
            if memory < VALIDATOR_MIN_MEMORY_MB:
                logger.warning(f"Node {node_id[:8]}... cannot be Validator: insufficient memory ({memory}MB)")
                return False
            
            # Check if already a validator
            last_layer = max(0, self.current_num_layers - 1)
            if any(a.node_id == node_id for a in self.layer_assignments[last_layer]):
                logger.info(f"Node {node_id[:8]}... is already a Validator")
                return True
            
            # Assign the last layer (LM head)
            self._assign_layer(last_layer, node_id, node_url, grpc_addr)
            self.lm_head_holder = node_id
            
            logger.info(f"Node {node_id[:8]}... upgraded to VALIDATOR (assigned layer {last_layer})")
            return True
    
    def demote_from_validator(self, node_id: str) -> bool:
        """
        Demote a node from Validator role when stake drops below requirement.
        
        This is called when:
        1. A validator unstakes and drops below the required amount
        2. The network grows and the required stake increases (tier change)
        
        The node keeps its other layer assignments but loses the LM head.
        
        Returns True if demotion was successful.
        """
        with self.lock:
            return self._demote_from_validator_unlocked(node_id)
    
    def _demote_from_validator_unlocked(self, node_id: str) -> bool:
        """
        Internal demotion logic (caller must hold self.lock).
        
        Split out to avoid deadlock when called from validate_all_validators().
        """
        last_layer = max(0, self.current_num_layers - 1)
        
        # Check if node is currently a validator
        current_assignments = self.layer_assignments.get(last_layer, [])
        was_validator = any(a.node_id == node_id for a in current_assignments)
        
        if not was_validator:
            logger.debug(f"Node {node_id[:8]}... is not a validator, nothing to demote")
            return False
        
        # Remove from last layer assignments
        self.layer_assignments[last_layer] = [
            a for a in current_assignments if a.node_id != node_id
        ]
        
        # Update lm_head_holder if this was the holder
        if self.lm_head_holder == node_id:
            # Find another validator if available
            remaining = self.layer_assignments.get(last_layer, [])
            if remaining:
                self.lm_head_holder = remaining[0].node_id
            else:
                self.lm_head_holder = None
        
        logger.warning(f"Node {node_id[:8]}... DEMOTED from Validator (insufficient stake)")
        return True
    
    def validate_all_validators(self, get_stake_fn) -> List[str]:
        """
        Validate all current validators still meet stake requirements.
        
        Called periodically or when stake tier changes to ensure all validators
        have sufficient stake for the current network size.
        
        IMPORTANT: Never demotes below MIN_VALIDATORS (2) to ensure the network
        can always compute real loss. The stake requirement only applies to
        validators beyond the minimum.
        
        Args:
            get_stake_fn: Function(node_id) -> float that returns current stake
            
        Returns:
            List of node_ids that were demoted
        """
        from neuroshard.core.economics.constants import get_dynamic_validator_stake
        
        MIN_VALIDATORS = 2  # Network needs at least 2 validators to function
        
        demoted = []
        
        with self.lock:
            last_layer = max(0, self.current_num_layers - 1)
            current_validators = list(self.layer_assignments.get(last_layer, []))
            num_validators = len(current_validators)
            
            # CRITICAL: Never demote below MIN_VALIDATORS
            # Otherwise the network can't compute real cross-entropy loss!
            if num_validators <= MIN_VALIDATORS:
                logger.debug(f"Only {num_validators} validators - skipping stake check (minimum {MIN_VALIDATORS} required)")
                return []
            
            # Get current stake requirement
            required_stake = get_dynamic_validator_stake(num_validators)
            
            # Sort validators by stake (lowest first) to demote lowest-stake first
            validators_with_stake = [
                (assignment, get_stake_fn(assignment.node_id))
                for assignment in current_validators
            ]
            validators_with_stake.sort(key=lambda x: x[1])  # Lowest stake first
            
            for assignment, node_stake in validators_with_stake:
                # Check if we'd go below minimum
                remaining_validators = num_validators - len(demoted)
                if remaining_validators <= MIN_VALIDATORS:
                    logger.info(f"Stopping demotion: {remaining_validators} validators remain (minimum {MIN_VALIDATORS})")
                    break
                
                if node_stake < required_stake:
                    logger.warning(
                        f"Validator {assignment.node_id[:8]}... has {node_stake:.0f} NEURO "
                        f"but {required_stake:.0f} required - DEMOTING"
                    )
                    # Use unlocked version since we already hold self.lock
                    if self._demote_from_validator_unlocked(assignment.node_id):
                        demoted.append(assignment.node_id)
        
        return demoted
    
    def unregister_node(self, node_id: str):
        """
        Unregister a node and redistribute its layers.
        
        This handles graceful degradation when nodes leave.
        """
        with self.lock:
            # Remove from capacities
            self.node_capacities.pop(node_id, None)
            
            # Find all layers this node was holding
            orphaned_layers = []
            
            for layer_id, assignments in self.layer_assignments.items():
                # Remove this node's assignment
                self.layer_assignments[layer_id] = [
                    a for a in assignments if a.node_id != node_id
                ]
                
                # Check if layer is now orphaned (< MIN_REPLICAS)
                if len(self.layer_assignments[layer_id]) < self.MIN_REPLICAS:
                    orphaned_layers.append(layer_id)
            
            # Handle embedding/head holder leaving
            if self.embedding_holder == node_id:
                self.embedding_holder = None
            if self.lm_head_holder == node_id:
                self.lm_head_holder = None
            
            if orphaned_layers:
                logger.warning(f"Node {node_id[:8]}... left, {len(orphaned_layers)} layers need redistribution")
                # In production, we would trigger redistribution here
    
    def get_layer_holders(self, layer_id: int) -> List[LayerAssignment]:
        """Get all nodes holding a specific layer."""
        with self.lock:
            return list(self.layer_assignments.get(layer_id, []))
    
    def get_pipeline_route(self) -> List[Tuple[int, str]]:
        """
        Get the route for pipeline inference.
        
        Returns list of (layer_id, grpc_addr) for each layer in order.
        
        Filters out dead/stale nodes based on heartbeat timeout.
        """
        with self.lock:
            route = []
            now = time.time()
            
            for layer_id in range(self.current_num_layers):
                holders = self.layer_assignments.get(layer_id, [])
                if not holders:
                    logger.error(f"Layer {layer_id} has no holders!")
                    continue
                
                # ROBUSTNESS: Filter out stale holders (expired heartbeat)
                active_holders = [
                    h for h in holders
                    if (now - h.last_heartbeat) < self.HEARTBEAT_TIMEOUT
                ]
                
                if not active_holders:
                    logger.warning(f"Layer {layer_id} has no ACTIVE holders "
                                  f"({len(holders)} total, all stale)")
                    continue
                
                # Pick best active holder (most recent heartbeat)
                active_holders.sort(key=lambda h: -h.last_heartbeat)
                route.append((layer_id, active_holders[0].grpc_addr))
                
                logger.debug(f"Layer {layer_id}: selected {active_holders[0].node_id[:16]}... "
                            f"(heartbeat {now - active_holders[0].last_heartbeat:.1f}s ago)")
            
            return route
    
    def get_network_capacity(self) -> NetworkCapacity:
        """Get current network capacity with dynamic architecture."""
        with self.lock:
            total_memory = sum(self.node_capacities.values())
            
            # Calculate max layers based on current architecture
            if self.current_architecture:
                memory_per_layer = estimate_memory_per_layer(self.current_architecture)
                max_layers = int(total_memory * 0.6 / memory_per_layer)
            else:
                max_layers = 0
            
            layer_coverage = {
                layer_id: len(assignments)
                for layer_id, assignments in self.layer_assignments.items()
            }
            
            return NetworkCapacity(
                total_nodes=len(self.node_capacities),
                total_memory_mb=total_memory,
                max_layers=max_layers,
                assigned_layers=self.current_num_layers,
                layer_coverage=layer_coverage,
            )
    
    def heartbeat(self, node_id: str, layer_ids: List[int]):
        """Update heartbeat for a node's layers."""
        with self.lock:
            now = time.time()
            for layer_id in layer_ids:
                for assignment in self.layer_assignments.get(layer_id, []):
                    if assignment.node_id == node_id:
                        assignment.last_heartbeat = now
    
    def cleanup_stale_assignments(self) -> int:
        """
        Remove stale layer assignments (nodes that haven't heartbeat recently).
        
        Returns number of stale assignments removed.
        
        Called periodically to prevent dead peers from being selected for pipeline routing.
        """
        with self.lock:
            now = time.time()
            removed_count = 0
            
            for layer_id, assignments in list(self.layer_assignments.items()):
                # Filter out stale assignments
                active_assignments = [
                    a for a in assignments
                    if (now - a.last_heartbeat) < self.HEARTBEAT_TIMEOUT
                ]
                
                stale_count = len(assignments) - len(active_assignments)
                if stale_count > 0:
                    logger.info(f"Layer {layer_id}: removed {stale_count} stale assignments "
                               f"({len(active_assignments)} remain)")
                    removed_count += stale_count
                
                # Update assignments
                if active_assignments:
                    self.layer_assignments[layer_id] = active_assignments
                else:
                    # No active holders for this layer!
                    logger.warning(f"Layer {layer_id}: NO active holders remaining!")
                    del self.layer_assignments[layer_id]
            
            return removed_count


class DynamicNeuroLLM:
    """
    A NeuroLLM that dynamically scales with the network.
    
    Key differences from fixed-phase model:
    - Number of layers AND hidden dimension determined by network capacity
    - Layers are distributed across nodes
    - Model grows organically in BOTH width and depth
    - Architecture adapts automatically as network expands
    """
    
    def __init__(
        self,
        node_id: str,
        layer_pool: DynamicLayerPool,
        device: str = "cpu"
    ):
        self.node_id = node_id
        self.layer_pool = layer_pool
        self.device = device
        
        # Get current architecture from layer pool
        if layer_pool.current_architecture is None:
            raise RuntimeError("Layer pool has no architecture - call _auto_recalculate_architecture first")
        self.architecture = layer_pool.current_architecture
        
        # My assigned layers
        self.my_layers: Dict[int, torch.nn.Module] = {}
        self.my_layer_ids: List[int] = []
        
        # Callback for when layers change (set by DynamicNeuroNode to sync state)
        self._on_layers_changed: Optional[callable] = None
        
        # Reference to P2P manager for DHT updates during layer removal
        self._p2p_manager = None
        
        # Do I hold embedding/head?
        self.has_embedding = False
        self.has_lm_head = False
        
        # Shared components (if I hold them)
        self.embedding: Optional[torch.nn.Embedding] = None
        self.lm_head: Optional[torch.nn.Linear] = None
        self.final_norm: Optional[torch.nn.Module] = None
        
        # Training mode flag (PyTorch convention)
        self.training = False
        
        logger.info(f"DynamicNeuroLLM initialized for node {node_id[:8]}... "
                   f"with {self.architecture.num_layers}L × {self.architecture.hidden_dim}H architecture")
    
    def initialize_layers(self, layer_ids: List[int]):
        """Initialize the layers assigned to this node using DYNAMIC architecture."""
        from neuroshard.core.model.llm import NeuroLLMConfig, NeuroDecoderLayer
        
        # Create config from current architecture (DYNAMIC!)
        config = NeuroLLMConfig(
            hidden_dim=self.architecture.hidden_dim,
            intermediate_dim=self.architecture.intermediate_dim,
            num_layers=self.architecture.num_layers,
            num_heads=self.architecture.num_heads,
            num_kv_heads=self.architecture.num_kv_heads,
            vocab_size=self.architecture.vocab_size,
            max_seq_len=self.architecture.max_seq_len,
            dropout=self.architecture.dropout,
            rope_theta=self.architecture.rope_theta,
        )
        
        for layer_id in layer_ids:
            layer = NeuroDecoderLayer(config, layer_id)
            layer.to(self.device)
            self.my_layers[layer_id] = layer
        
        self.my_layer_ids = sorted(layer_ids)
        
        # Initialize embedding if I'm the holder (uses dynamic hidden_dim!)
        # Vocab capacity can grow dynamically as tokenizer learns more merges
        self.vocab_capacity = INITIAL_VOCAB_SIZE
        if self.layer_pool.embedding_holder == self.node_id:
            self.embedding = torch.nn.Embedding(self.vocab_capacity, self.architecture.hidden_dim)
            self.embedding.to(self.device)
            self.has_embedding = True
        
        # Initialize LM head if I'm the holder (uses dynamic hidden_dim!)
        if self.layer_pool.lm_head_holder == self.node_id:
            self.lm_head = torch.nn.Linear(self.architecture.hidden_dim, self.vocab_capacity, bias=False)
            from neuroshard.core.model.llm import RMSNorm
            self.final_norm = RMSNorm(self.architecture.hidden_dim)
            self.lm_head.to(self.device)
            self.final_norm.to(self.device)
            self.has_lm_head = True
        
        logger.info(f"Initialized {len(layer_ids)} layers: {layer_ids}, "
                   f"arch={self.architecture.num_layers}L×{self.architecture.hidden_dim}H, "
                   f"embedding={self.has_embedding}, head={self.has_lm_head}")
    
    def initialize_lm_head(self) -> bool:
        """
        Dynamically initialize the LM head (for validator upgrade).
        
        Called when a node is upgraded to Validator after staking.
        No restart required - initializes the head in place.
        
        Returns True if initialization was successful.
        """
        if self.has_lm_head:
            logger.info("LM head already initialized")
            return True
        
        try:
            from neuroshard.core.model.llm import RMSNorm
            
            self.lm_head = torch.nn.Linear(self.architecture.hidden_dim, self.vocab_capacity, bias=False)
            self.final_norm = RMSNorm(self.architecture.hidden_dim)
            self.lm_head.to(self.device)
            self.final_norm.to(self.device)
            self.has_lm_head = True
            
            # Add last layer to my layers if not already there
            last_layer = self.architecture.num_layers - 1
            if last_layer not in self.my_layer_ids:
                self.my_layer_ids.append(last_layer)
                self.my_layer_ids = sorted(self.my_layer_ids)
            
            logger.info(f"LM head initialized! Now computing REAL cross-entropy loss")
            return True
        except Exception as e:
            logger.error(f"Failed to initialize LM head: {e}")
    
    def disable_lm_head(self) -> bool:
        """
        Disable the LM head (for validator demotion).
        
        Called when a validator is demoted due to insufficient stake.
        The node reverts to Worker role and uses activation norm as loss.
        
        Returns True if demotion was successful.
        """
        if not self.has_lm_head:
            logger.debug("LM head not initialized, nothing to disable")
            return False
        
        try:
            # Free memory from LM head
            if self.lm_head is not None:
                del self.lm_head
                self.lm_head = None
            if self.final_norm is not None:
                del self.final_norm
                self.final_norm = None
            
            self.has_lm_head = False
            
            # Force garbage collection to free memory
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.warning(f"LM head DISABLED - node demoted to Worker (will use activation norm as loss)")
            return True
        except Exception as e:
            logger.error(f"Failed to disable LM head: {e}")
            return False
    
    def add_layers(self, new_layer_ids: List[int]) -> List[int]:
        """
        Dynamically add new layers to a running model.
        
        This allows the model to grow during training without restart!
        Called when:
        - Network needs more layers
        - Node has available memory
        - Vocab expansion freed up memory
        
        Args:
            new_layer_ids: Layer IDs to add
            
        Returns:
            List of layer IDs that were successfully added
        """
        from neuroshard.core.model.llm import NeuroLLMConfig, NeuroDecoderLayer
        
        # Filter out layers we already have
        layers_to_add = [lid for lid in new_layer_ids if lid not in self.my_layers]
        if not layers_to_add:
            return []
        
        # Check memory before adding
        memory_per_layer = estimate_memory_per_layer(self.architecture)
        required_mb = memory_per_layer * len(layers_to_add)
        
        try:
            if torch.cuda.is_available() and self.device != "cpu":
                free_mb = (torch.cuda.get_device_properties(0).total_memory - 
                          torch.cuda.memory_allocated()) / (1024 * 1024)
            else:
                import psutil
                free_mb = psutil.virtual_memory().available / (1024 * 1024)
            
            if free_mb < required_mb * 1.5:  # 1.5x safety margin
                logger.warning(f"[LAYER] Insufficient memory to add {len(layers_to_add)} layers: "
                              f"need {required_mb:.0f}MB, have {free_mb:.0f}MB")
                # Add as many as we can fit
                can_fit = int(free_mb / (memory_per_layer * 1.5))
                layers_to_add = layers_to_add[:can_fit]
                if not layers_to_add:
                    return []
        except Exception as e:
            logger.warning(f"[LAYER] Memory check failed: {e}, proceeding cautiously")
        
        # Create config from architecture
        config = NeuroLLMConfig(
            hidden_dim=self.architecture.hidden_dim,
            intermediate_dim=self.architecture.intermediate_dim,
            num_layers=self.architecture.num_layers,
            num_heads=self.architecture.num_heads,
            num_kv_heads=self.architecture.num_kv_heads,
            vocab_size=self.architecture.vocab_size,
            max_seq_len=self.architecture.max_seq_len,
            dropout=self.architecture.dropout,
            rope_theta=self.architecture.rope_theta,
        )
        
        added = []
        for layer_id in layers_to_add:
            try:
                layer = NeuroDecoderLayer(config, layer_id)
                layer.to(self.device)
                self.my_layers[layer_id] = layer
                added.append(layer_id)
            except Exception as e:
                logger.error(f"[LAYER] Failed to add layer {layer_id}: {e}")
                break  # Stop on first failure
        
        if added:
            self.my_layer_ids = sorted(self.my_layers.keys())
            logger.info(f"[LAYER] ✅ Added {len(added)} layers: {added}")
            logger.info(f"[LAYER] Now holding {len(self.my_layer_ids)} layers: {self.my_layer_ids[:5]}...{self.my_layer_ids[-5:]}")
        
        return added
    
    def remove_layers(self, layer_ids_to_remove: List[int], layer_pool=None, p2p_manager=None) -> List[int]:
        """
        Dynamically remove layers from a running model.
        
        This allows the model to shrink during training without restart!
        Called when:
        - Vocab growth needs more memory
        - Network is redistributing layers
        - Memory pressure detected
        
        IMPORTANT: This also updates layer_pool and DHT announcements so other
        nodes know we no longer hold these layers!
        
        Args:
            layer_ids_to_remove: Layer IDs to remove
            layer_pool: DynamicLayerPool to update assignments (optional)
            p2p_manager: P2PManager to update DHT announcements (optional)
            
        Returns:
            List of layer IDs that were successfully removed
        """
        # Can't remove layers we don't have
        layers_to_remove = [lid for lid in layer_ids_to_remove if lid in self.my_layers]
        if not layers_to_remove:
            return []
        
        # Don't remove all layers - keep at least 1
        if len(layers_to_remove) >= len(self.my_layers):
            layers_to_remove = layers_to_remove[:-1]  # Keep last one
            if not layers_to_remove:
                logger.warning("[LAYER] Cannot remove all layers")
                return []
        
        removed = []
        for layer_id in layers_to_remove:
            try:
                # Delete the layer
                layer = self.my_layers.pop(layer_id)
                del layer
                removed.append(layer_id)
            except Exception as e:
                logger.error(f"[LAYER] Failed to remove layer {layer_id}: {e}")
        
        if removed:
            self.my_layer_ids = sorted(self.my_layers.keys())
            
            # UPDATE LAYER POOL: Remove this node from removed layers' assignments
            if layer_pool:
                try:
                    with layer_pool.lock:
                        for layer_id in removed:
                            if layer_id in layer_pool.layer_assignments:
                                # Remove our node from this layer's holders
                                layer_pool.layer_assignments[layer_id] = [
                                    a for a in layer_pool.layer_assignments[layer_id]
                                    if a.node_id != self.node_id
                                ]
                                logger.info(f"[LAYER] Updated layer_pool: removed self from layer {layer_id}")
                except Exception as e:
                    logger.warning(f"[LAYER] Could not update layer_pool: {e}")
            
            # UPDATE DHT: Change announced layer range
            if p2p_manager and self.my_layer_ids:
                try:
                    new_start = min(self.my_layer_ids)
                    new_end = max(self.my_layer_ids)
                    p2p_manager.start_layer = new_start
                    p2p_manager.end_layer = new_end
                    p2p_manager.shard_range = f"{new_start}-{new_end}"
                    logger.info(f"[LAYER] Updated P2P shard_range: {new_start}-{new_end}")
                    
                    # Re-announce immediately so network knows
                    if hasattr(p2p_manager, '_announce_once'):
                        p2p_manager._announce_once(verbose=True)
                except Exception as e:
                    logger.warning(f"[LAYER] Could not update P2P shard_range: {e}")
            
            # Force garbage collection to free memory
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.info(f"[LAYER] ✅ Removed {len(removed)} layers: {removed}")
            logger.info(f"[LAYER] Now holding {len(self.my_layer_ids)} layers")
            
            # NOTIFY: Call callback so node can sync its state
            if self._on_layers_changed:
                try:
                    self._on_layers_changed(self.my_layer_ids)
                except Exception as e:
                    logger.warning(f"[LAYER] Callback failed: {e}")
        
        return removed
    
    def expand_vocabulary(self, new_vocab_size: int) -> bool:
        """
        Expand embedding and lm_head to accommodate a larger vocabulary.
        
        This is called when the tokenizer learns new BPE merges that exceed
        the current vocabulary capacity. The expansion preserves existing
        token embeddings while initializing new ones.
        
        For an ever-growing decentralized LLM, vocabulary expansion is essential
        as millions of users contribute diverse training data across languages
        and domains.
        
        Args:
            new_vocab_size: The new vocabulary size (must be > current capacity)
            
        Returns:
            True if expansion was successful, False otherwise
        """
        if new_vocab_size <= self.vocab_capacity:
            return True  # No expansion needed
        
        # Check against max (if set)
        if MAX_VOCAB_SIZE is not None and new_vocab_size > MAX_VOCAB_SIZE:
            logger.warning(f"Requested vocab {new_vocab_size} exceeds MAX_VOCAB_SIZE {MAX_VOCAB_SIZE}")
            new_vocab_size = MAX_VOCAB_SIZE
        
        # Round up to next VOCAB_GROWTH_CHUNK for efficient memory alignment
        new_capacity = ((new_vocab_size + VOCAB_GROWTH_CHUNK - 1) // VOCAB_GROWTH_CHUNK) * VOCAB_GROWTH_CHUNK
        if MAX_VOCAB_SIZE is not None:
            new_capacity = min(new_capacity, MAX_VOCAB_SIZE)
        
        # MEMORY CHECK: Estimate memory needed for expansion
        # Embedding + LM head expansion: 2 * (new - old) * hidden_dim * 4 bytes (float32)
        hidden_dim = self.architecture.hidden_dim
        expansion_params = 2 * (new_capacity - self.vocab_capacity) * hidden_dim
        expansion_memory_mb = (expansion_params * 4) / (1024 * 1024)  # Just weights, no optimizer yet
        
        # Check available memory (GPU or CPU)
        try:
            if torch.cuda.is_available() and self.device != "cpu":
                # Check GPU memory
                free_memory_mb = (torch.cuda.get_device_properties(0).total_memory - 
                                 torch.cuda.memory_allocated()) / (1024 * 1024)
            else:
                # Check system RAM
                import psutil
                free_memory_mb = psutil.virtual_memory().available / (1024 * 1024)
            
            # Need at least 2x expansion memory (weights + temporary copy during expansion)
            required_mb = expansion_memory_mb * 2
            
            if free_memory_mb < required_mb:
                logger.warning(f"[VOCAB] Insufficient memory for expansion: need {required_mb:.0f}MB, "
                              f"have {free_memory_mb:.0f}MB free")
                
                # TRY: Remove some layers to make room for vocab expansion
                # Vocab is more important than extra layers (all nodes need same vocab)
                memory_per_layer = estimate_memory_per_layer(self.architecture)
                layers_to_free = int((required_mb - free_memory_mb) / memory_per_layer) + 1
                
                if len(self.my_layers) > layers_to_free + 1:  # Keep at least 1 layer
                    # Remove highest-numbered layers (least important for Driver/Validator)
                    layers_to_remove = sorted(self.my_layers.keys(), reverse=True)[:layers_to_free]
                    logger.warning(f"[VOCAB] Attempting to free memory by removing {layers_to_free} layers: {layers_to_remove}")
                    
                    # Pass layer_pool so network state is updated
                    # p2p_manager will be notified via layer_pool.on_layers_changed callback
                    removed = self.remove_layers(
                        layers_to_remove, 
                        layer_pool=self.layer_pool,
                        p2p_manager=getattr(self, '_p2p_manager', None)
                    )
                    if removed:
                        # Recalculate free memory after layer removal
                        if torch.cuda.is_available() and self.device != "cpu":
                            free_memory_mb = (torch.cuda.get_device_properties(0).total_memory - 
                                             torch.cuda.memory_allocated()) / (1024 * 1024)
                        else:
                            import psutil
                            free_memory_mb = psutil.virtual_memory().available / (1024 * 1024)
                        
                        if free_memory_mb >= required_mb:
                            logger.info(f"[VOCAB] ✅ Freed enough memory by removing {len(removed)} layers")
                            # Continue with expansion below
                        else:
                            logger.warning(f"[VOCAB] Still insufficient memory after layer removal")
                            return False
                    else:
                        logger.warning(f"[VOCAB] Could not remove layers, capping expansion")
                        return False
                else:
                    logger.warning(f"[VOCAB] Not enough layers to remove, capping expansion")
                    return False
                
            logger.info(f"[VOCAB] Memory check passed: {expansion_memory_mb:.0f}MB needed, "
                       f"{free_memory_mb:.0f}MB available")
        except Exception as e:
            logger.warning(f"[VOCAB] Could not check memory: {e}, proceeding with expansion")
        
        logger.info(f"[VOCAB] Expanding vocabulary: {self.vocab_capacity} → {new_capacity}")
        
        try:
            # Expand embedding if we have it
            if self.has_embedding and self.embedding is not None:
                old_embedding = self.embedding
                old_vocab = old_embedding.weight.shape[0]
                hidden_dim = old_embedding.weight.shape[1]
                
                # Create new larger embedding
                new_embedding = torch.nn.Embedding(new_capacity, hidden_dim)
                new_embedding.to(self.device)
                
                # Copy existing embeddings
                with torch.no_grad():
                    new_embedding.weight[:old_vocab] = old_embedding.weight
                    # Initialize new embeddings with small random values
                    # (similar to how transformers initialize)
                    std = 0.02
                    new_embedding.weight[old_vocab:].normal_(mean=0.0, std=std)
                
                # Replace old embedding
                del self.embedding
                self.embedding = new_embedding
                
                logger.info(f"[VOCAB] Expanded embedding: {old_vocab} → {new_capacity} tokens")
            
            # Expand lm_head if we have it
            if self.has_lm_head and self.lm_head is not None:
                old_lm_head = self.lm_head
                old_vocab = old_lm_head.weight.shape[0]
                hidden_dim = old_lm_head.weight.shape[1]
                
                # Create new larger lm_head
                new_lm_head = torch.nn.Linear(hidden_dim, new_capacity, bias=False)
                new_lm_head.to(self.device)
                
                # Copy existing weights
                with torch.no_grad():
                    new_lm_head.weight[:old_vocab] = old_lm_head.weight
                    # Initialize new output weights
                    std = 0.02
                    new_lm_head.weight[old_vocab:].normal_(mean=0.0, std=std)
                
                # Replace old lm_head
                del self.lm_head
                self.lm_head = new_lm_head
                
                logger.info(f"[VOCAB] Expanded lm_head: {old_vocab} → {new_capacity} output classes")
            
            # Update capacity
            old_capacity = self.vocab_capacity
            self.vocab_capacity = new_capacity
            
            # Force garbage collection
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.info(f"[VOCAB] ✅ Vocabulary expansion complete: {old_capacity} → {new_capacity}")
            return True
            
        except Exception as e:
            logger.error(f"[VOCAB] Failed to expand vocabulary: {e}")
            return False
    
    def check_and_expand_vocab_if_needed(self) -> bool:
        """
        Check if tokenizer vocabulary exceeds current capacity and expand if needed.
        
        This should be called periodically (e.g., after loading new tokenizer from CDN)
        to ensure the model can handle all tokens in the current vocabulary.
        
        Returns:
            True if no expansion needed or expansion successful, False on failure
        """
        if self.tokenizer is None:
            return True
        
        current_vocab = self.tokenizer.current_vocab_size
        if current_vocab > self.vocab_capacity:
            logger.info(f"[VOCAB] Tokenizer vocab ({current_vocab}) exceeds capacity ({self.vocab_capacity})")
            return self.expand_vocabulary(current_vocab)
        
        return True
    
    def forward_my_layers(
        self,
        hidden_states: torch.Tensor,
        start_layer: Optional[int] = None,
        end_layer: Optional[int] = None,
    ) -> torch.Tensor:
        """Forward through my assigned layers."""
        if start_layer is None:
            start_layer = min(self.my_layer_ids) if self.my_layer_ids else 0
        if end_layer is None:
            end_layer = max(self.my_layer_ids) + 1 if self.my_layer_ids else 0
        
        x = hidden_states
        
        for layer_id in range(start_layer, end_layer):
            if layer_id in self.my_layers:
                x, _ = self.my_layers[layer_id](x)
        
        return x
    
    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Embed input tokens (only if I hold embedding)."""
        if not self.has_embedding:
            raise RuntimeError("This node does not hold the embedding layer")
        return self.embedding(input_ids)
    
    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Compute logits (only if I hold LM head)."""
        if not self.has_lm_head:
            raise RuntimeError("This node does not hold the LM head")
        x = self.final_norm(hidden_states)
        return self.lm_head(x)
    
    def get_num_params(self) -> int:
        """Get number of parameters on this node."""
        total = 0
        for layer in self.my_layers.values():
            total += sum(p.numel() for p in layer.parameters())
        if self.embedding:
            total += sum(p.numel() for p in self.embedding.parameters())
        if self.lm_head:
            total += sum(p.numel() for p in self.lm_head.parameters())
        if self.final_norm:
            total += sum(p.numel() for p in self.final_norm.parameters())
        return total
    
    def parameters(self):
        """Yield all parameters for this node's model components (for optimizer/gradient clipping)."""
        for layer in self.my_layers.values():
            yield from layer.parameters()
        if self.embedding:
            yield from self.embedding.parameters()
        if self.lm_head:
            yield from self.lm_head.parameters()
        if self.final_norm:
            yield from self.final_norm.parameters()
    
    def named_parameters(self, prefix: str = '', recurse: bool = True):
        """
        Yield (name, param) tuples for all parameters.
        
        This is the standard PyTorch interface for iterating over named parameters.
        """
        # Layers
        for layer_id, layer in self.my_layers.items():
            layer_prefix = f"{prefix}layers.{layer_id}." if prefix else f"layers.{layer_id}."
            for name, param in layer.named_parameters(prefix='', recurse=recurse):
                yield layer_prefix + name, param
        
        # Embedding
        if self.embedding:
            emb_prefix = f"{prefix}embedding." if prefix else "embedding."
            for name, param in self.embedding.named_parameters(prefix='', recurse=recurse):
                yield emb_prefix + name, param
        
        # LM Head
        if self.lm_head:
            head_prefix = f"{prefix}lm_head." if prefix else "lm_head."
            for name, param in self.lm_head.named_parameters(prefix='', recurse=recurse):
                yield head_prefix + name, param
        
        # Final Norm
        if self.final_norm:
            norm_prefix = f"{prefix}final_norm." if prefix else "final_norm."
            for name, param in self.final_norm.named_parameters(prefix='', recurse=recurse):
                yield norm_prefix + name, param
    
    def state_dict(self) -> Dict[str, Any]:
        """
        Return the state dictionary of the model.
        
        This is the standard PyTorch interface for saving model state.
        """
        state = {}
        
        # Layers
        for layer_id, layer in self.my_layers.items():
            for name, param in layer.state_dict().items():
                state[f"layers.{layer_id}.{name}"] = param
        
        # Embedding
        if self.embedding:
            for name, param in self.embedding.state_dict().items():
                state[f"embedding.{name}"] = param
        
        # LM Head
        if self.lm_head:
            for name, param in self.lm_head.state_dict().items():
                state[f"lm_head.{name}"] = param
        
        # Final Norm
        if self.final_norm:
            for name, param in self.final_norm.state_dict().items():
                state[f"final_norm.{name}"] = param
        
        return state
    
    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True):
        """
        Load state dictionary into the model.
        
        This is the standard PyTorch interface for loading model state.
        """
        # Group state by component
        layer_states: Dict[int, Dict[str, Any]] = {}
        embedding_state: Dict[str, Any] = {}
        lm_head_state: Dict[str, Any] = {}
        final_norm_state: Dict[str, Any] = {}
        
        for key, value in state_dict.items():
            if key.startswith("layers."):
                parts = key.split(".", 2)
                layer_id = int(parts[1])
                param_name = parts[2]
                if layer_id not in layer_states:
                    layer_states[layer_id] = {}
                layer_states[layer_id][param_name] = value
            elif key.startswith("embedding."):
                param_name = key[len("embedding."):]
                embedding_state[param_name] = value
            elif key.startswith("lm_head."):
                param_name = key[len("lm_head."):]
                lm_head_state[param_name] = value
            elif key.startswith("final_norm."):
                param_name = key[len("final_norm."):]
                final_norm_state[param_name] = value
        
        # Load into components
        for layer_id, layer in self.my_layers.items():
            if layer_id in layer_states:
                layer.load_state_dict(layer_states[layer_id], strict=strict)
        
        if self.embedding and embedding_state:
            self.embedding.load_state_dict(embedding_state, strict=strict)
        
        if self.lm_head and lm_head_state:
            self.lm_head.load_state_dict(lm_head_state, strict=strict)
        
        if self.final_norm and final_norm_state:
            self.final_norm.load_state_dict(final_norm_state, strict=strict)
    
    def zero_grad(self, set_to_none: bool = False):
        """
        Zero all gradients.
        
        This is the standard PyTorch interface for zeroing gradients.
        """
        for param in self.parameters():
            if param.grad is not None:
                if set_to_none:
                    param.grad = None
                else:
                    param.grad.zero_()
    
    def train(self, mode: bool = True) -> 'DynamicNeuroLLM':
        """
        Set the model to training mode.
        
        This is the standard PyTorch interface for setting training mode
        on all submodules.
        """
        self.training = mode
        for layer in self.my_layers.values():
            layer.train(mode)
        if self.embedding:
            self.embedding.train(mode)
        if self.lm_head:
            self.lm_head.train(mode)
        if self.final_norm:
            self.final_norm.train(mode)
        return self
    
    def eval(self) -> 'DynamicNeuroLLM':
        """Set the model to evaluation mode."""
        return self.train(False)
    
    def get_my_contribution(self) -> Dict[str, Any]:
        """Get this node's contribution to the network."""
        capacity = self.layer_pool.get_network_capacity()
        
        return {
            "node_id": self.node_id[:16] + "...",
            "my_layers": self.my_layer_ids,
            "my_params": self.get_num_params(),
            "has_embedding": self.has_embedding,
            "has_lm_head": self.has_lm_head,
            "network_total_layers": capacity.assigned_layers,
            "network_total_nodes": capacity.total_nodes,
            "contribution_ratio": len(self.my_layer_ids) / max(1, capacity.assigned_layers),
        }


def calculate_reward_multiplier(
    num_layers_held: int,
    total_network_layers: int,
    has_embedding: bool,
    has_lm_head: bool
) -> float:
    """
    Calculate NEURO reward multiplier based on contribution.
    
    Roles:
    - Worker: Standard reward based on layers
    - Driver (Embedding): 1.2x bonus (bandwidth cost)
    - Validator (Head): 1.2x bonus (compute/consensus cost)
    """
    if total_network_layers == 0:
        return 1.0
    
    # Base multiplier from layer contribution
    layer_ratio = num_layers_held / total_network_layers
    base_multiplier = 1.0 + layer_ratio  # 1.0 to 2.0 based on layers
    
    # Bonus for critical components (Roles)
    if has_embedding:
        base_multiplier *= 1.2  # 20% bonus for Driving (Data bandwidth)
    if has_lm_head:
        base_multiplier *= 1.2  # 20% bonus for Validating (Loss calc + Gradient origin)
    
    return base_multiplier


# ============================================================================
# DYNAMIC NEURO NODE - The Main Node Class
# ============================================================================

class DynamicNeuroNode:
    """
    A truly decentralized node that contributes based on available memory.
    
    NO PHASES. NO CENTRAL COORDINATION.
    
    How it works:
    1. Node starts, detects available memory
    2. Registers with network, gets assigned layers
    3. Loads only the layers it's responsible for
    4. Participates in training (computes gradients for its layers)
    5. Participates in inference (forwards through its layers)
    6. Earns NEURO proportional to its contribution
    
    The more memory you have, the more layers you hold, the more you earn.
    """
    
    CHECKPOINT_DIR = None  # Set in __init__
    
    def __init__(
        self,
        node_id: str,
        port: int = 8000,
        tracker_url: str = "https://neuroshard.com/api/tracker",
        node_token: Optional[str] = None,
        device: str = "cpu",
        available_memory_mb: Optional[float] = None,
        enable_training: bool = True,
        max_storage_mb: float = 100.0,
        max_cpu_threads: Optional[int] = None,
    ):
        self.node_id = node_id
        self.port = port
        self.tracker_url = tracker_url
        self.node_token = node_token
        
        # Detect device automatically if "auto" or "cpu" (backward compatibility)
        if device in ("auto", "cpu"):
            if torch.cuda.is_available():
                self.device = "cuda"
                logger.info(f"[NODE] GPU detected: CUDA available")
            elif torch.backends.mps.is_available():
                # MPS (Apple Silicon GPU) - enabled now that we have GIL yields
                self.device = "mps"
                logger.info(f"[NODE] GPU detected: Apple Metal (MPS)")
            else:
                self.device = "cpu"
                logger.info(f"[NODE] No GPU detected, using CPU")
                
                # Help debug why CUDA isn't available
                import subprocess
                import sys
                
                # Check if NVIDIA GPU exists
                has_nvidia_gpu = False
                try:
                    if sys.platform == 'win32':
                        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True, timeout=2)
                        has_nvidia_gpu = result.returncode == 0
                    elif sys.platform == 'darwin':
                        # macOS doesn't have NVIDIA support (use MPS instead)
                        pass
                    else:  # Linux
                        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True, timeout=2)
                        has_nvidia_gpu = result.returncode == 0
                except Exception:
                    pass
                
                # Detailed diagnostic
                logger.info(f"[NODE] torch.cuda.is_available() = False")
                
                # Check if PyTorch was built with CUDA
                try:
                    cuda_available = torch.cuda.is_available()
                    cuda_built = getattr(torch.version, 'cuda', None)
                    torch_version = torch.__version__
                    logger.info(f"[NODE] PyTorch version: {torch_version}")
                    logger.info(f"[NODE] CUDA compiled version: {cuda_built if cuda_built else 'None (CPU-only build)'}")
                except Exception as e:
                    logger.info(f"[NODE] Could not get CUDA info: {e}")
                
                # Provide helpful diagnostic
                if has_nvidia_gpu:
                    logger.warning("⚠️  NVIDIA GPU DETECTED BUT NOT BEING USED!")
                    logger.warning("Your system has an NVIDIA GPU, but this PyTorch installation is CPU-only.")
                    logger.warning("🔧 TO ENABLE GPU (for 5-10x faster training):")
                    logger.warning("If running the .exe (frozen build):")
                    logger.warning("  Unfortunately, the bundled Python environment can't easily be modified.")
                    logger.warning("  We recommend running from source for GPU support.")
                    logger.warning("If running from source:")
                    logger.warning("  pip uninstall torch")
                    logger.warning("  pip install torch --index-url https://download.pytorch.org/whl/cu121")
                    logger.warning("To verify: python -c \"import torch; print(torch.cuda.is_available())\"")
        else:
            self.device = device
            logger.info(f"[NODE] Device manually set to: {self.device}")
            
        logger.info(f"Using device: {self.device}")
        
        self.enable_training = enable_training
        self.max_storage_mb = max_storage_mb
        self.max_cpu_threads = max_cpu_threads
        
        # CPU thread limiting is done in runner.py BEFORE any torch operations
        # (torch.set_num_interop_threads must be called before any parallel work)
        if max_cpu_threads and self.device == "cpu":
            torch.set_num_threads(max_cpu_threads)  # Intra-op parallelism only
            logger.info(f"Set PyTorch intra-op threads: {max_cpu_threads}")
        
        # Detect memory if not provided
        if available_memory_mb is None:
            self.available_memory_mb = self._detect_available_memory()
        else:
            self.available_memory_mb = available_memory_mb
        
        # Checkpoint directory
        from pathlib import Path
        self.CHECKPOINT_DIR = Path.home() / ".neuroshard" / "checkpoints"
        self.CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
        
        # Compute wallet_id from token for stable checkpoint naming
        # (node_id can change if machine_id changes, but wallet_id is stable)
        if node_token:
            self.wallet_id = hashlib.sha256(node_token.encode()).hexdigest()[:16]
        else:
            self.wallet_id = self.node_id[:16]  # Fallback to node_id
        
        # Layer pool (shared across network via DHT)
        self.layer_pool: Optional[DynamicLayerPool] = None
        
        # My model (only my layers)
        self.model: Optional[DynamicNeuroLLM] = None
        self.my_layer_ids: List[int] = []
        
        # Tokenizer
        self.tokenizer = None
        
        # Training components (enable_training set in __init__)
        self.optimizer: Optional[torch.optim.Optimizer] = None
        self.training_coordinator = None
        self.data_manager = None
        self.gradient_gossip = None
        
        # Training lock to prevent concurrent training operations
        # (local training vs pipeline training conflict)
        self._training_lock = threading.Lock()
        
        # P2P
        self.p2p_manager = None
        
        # Stats
        self.is_running = False
        self.total_tokens_processed = 0
        self.total_training_rounds = 0
        self.current_loss = float('inf')
        self.inference_count = 0
        self.training_contribution_count = 0
        
        # KV cache for inference
        self.kv_cache: Dict[str, Any] = {}
        
        # Training context (keeps tensors alive for backward pass)
        # session_id -> {input, output, prev_peer}
        self.training_context: Dict[str, Any] = {}
        
        logger.info(f"DynamicNeuroNode initialized: memory={self.available_memory_mb:.0f}MB")
    
    def _detect_available_memory(self) -> float:
        """Detect available system memory."""
        try:
            import psutil
            mem = psutil.virtual_memory()
            # Use 70% of available memory for safety
            return mem.available * 0.7 / (1024 * 1024)
        except ImportError:
            # Fallback
            return 2000  # Assume 2GB
    
    def start(self):
        """Start the node."""
        logger.info("Starting DynamicNeuroNode...")
        
        # 1. Initialize layer pool
        dht = None
        if self.p2p_manager and hasattr(self.p2p_manager, 'dht'):
            dht = self.p2p_manager.dht
        self.layer_pool = DynamicLayerPool(dht_protocol=dht)
        
        # Pass device hint for memory calculations (CPU needs more conservative limits)
        self.layer_pool._device_hint = self.device
        
        # 1b. SMART ARCHITECTURE RECONCILIATION
        # This handles the case where the network has evolved while we were offline
        self._reconcile_architecture()
        
        # 1c. PRE-FETCH TOKENIZER VOCAB SIZE for accurate memory calculation
        # This is critical for dynamic vocab - we need to know vocab size BEFORE
        # assigning layers, otherwise we'll assign too many and OOM when vocab expands
        self._prefetch_vocab_capacity()
        
        # 2. Get staked amount from ledger (for Validator eligibility)
        staked_amount = 0.0
        if self.p2p_manager and self.p2p_manager.ledger:
            try:
                account_info = self.p2p_manager.ledger.get_account_info()
                staked_amount = account_info.get("stake", 0.0)
                logger.info(f"Current stake: {staked_amount:.2f} NEURO")
            except Exception as e:
                logger.debug(f"Could not get stake info: {e}")
        
        # 3. Register with network and get layer assignments
        self.my_layer_ids = self.layer_pool.register_node(
            node_id=self.node_id,
            node_url=f"http://localhost:{self.port}",
            grpc_addr=f"localhost:{self.port + 1000}",
            available_memory_mb=self.available_memory_mb,
            staked_amount=staked_amount
        )
        
        logger.info(f"Assigned {len(self.my_layer_ids)} layers: {self.my_layer_ids}")
        
        # 3. Initialize model with my layers
        self.model = DynamicNeuroLLM(
            node_id=self.node_id,
            layer_pool=self.layer_pool,
            device=self.device
        )
        self.model.initialize_layers(self.my_layer_ids)
        
        # 3b. Set up callback for dynamic layer changes
        # When model removes layers (e.g., for vocab expansion), sync node state
        def on_layers_changed(new_layer_ids: List[int]):
            self.my_layer_ids = new_layer_ids
            # Update P2P shard_range if available
            if self.p2p_manager and new_layer_ids:
                new_start = min(new_layer_ids)
                new_end = max(new_layer_ids)
                self.p2p_manager.start_layer = new_start
                self.p2p_manager.end_layer = new_end
                self.p2p_manager.shard_range = f"{new_start}-{new_end}"
                logger.info(f"[NODE] Synced layer_ids after change: {new_layer_ids}")
        
        self.model._on_layers_changed = on_layers_changed
        
        # 4. Initialize tokenizer with learned BPE merges from CDN
        from neuroshard.core.model.tokenizer import get_neuro_tokenizer, NeuroTokenizer
        self.tokenizer = get_neuro_tokenizer()
        self._load_learned_tokenizer()  # Update with BPE merges from CDN
        
        # 5. Try to load existing checkpoint (resume training)
        self._load_checkpoint()
        
        # 6. Setup training
        if self.enable_training:
            self._setup_training()
        
        self.is_running = True
        
        # Log contribution
        contribution = self.model.get_my_contribution()
        logger.info(f"Node started: {contribution['my_params']/1e6:.1f}M params, "
                   f"{len(self.my_layer_ids)} layers, "
                   f"embed={self.model.has_embedding}, head={self.model.has_lm_head}")
        
        # Verify model is actually on the expected device
        if self.my_layer_ids and self.model.my_layers:
            first_layer = self.model.my_layers[self.my_layer_ids[0]]
            param_device = next(first_layer.parameters()).device
            if str(param_device) != self.device and not (self.device == "cuda" and "cuda" in str(param_device)):
                logger.error(f"[DEVICE] Model device mismatch! Expected {self.device}, got {param_device}")
            else:
                logger.info(f"[DEVICE] Model verified on: {param_device}")
    
    def _load_learned_tokenizer(self):
        """
        Load learned BPE tokenizer from Genesis CDN.
        
        This ensures the tokenizer used for inference matches the one used
        for training data tokenization, providing consistency across the network.
        """
        import requests
        import os
        
        GENESIS_CDN_URL = "https://dwquwt9gkkeil.cloudfront.net"
        cache_dir = os.path.join(os.path.expanduser("~"), ".neuroshard", "data_cache")
        
        try:
            tokenizer_url = f"{GENESIS_CDN_URL}/tokenizer.json"
            tokenizer_cache_path = os.path.join(cache_dir, "tokenizer.json")
            
            # Try to fetch from CDN
            try:
                logger.debug(f"[TOKENIZER] Checking for learned tokenizer from {tokenizer_url}...")
                resp = requests.get(tokenizer_url, timeout=10)
                
                if resp.status_code == 200:
                    remote_tokenizer_data = resp.json()
                    remote_vocab_size = remote_tokenizer_data.get("next_merge_id", 0)
                    
                    # Cache locally
                    os.makedirs(cache_dir, exist_ok=True)
                    with open(tokenizer_cache_path, 'w') as f:
                        f.write(resp.text)
                    
                    # Update tokenizer if remote has more merges
                    if remote_vocab_size > self.tokenizer.next_merge_id:
                        from neuroshard.core.model.tokenizer import NeuroTokenizer
                        learned_tokenizer = NeuroTokenizer.load(tokenizer_cache_path)
                        
                        self.tokenizer.merges = learned_tokenizer.merges
                        self.tokenizer.merge_to_tokens = learned_tokenizer.merge_to_tokens
                        self.tokenizer.next_merge_id = learned_tokenizer.next_merge_id
                        
                        logger.info(f"[TOKENIZER] Loaded BPE tokenizer: {self.tokenizer.current_vocab_size} tokens, {len(self.tokenizer.merges)} merges")
                        
                        # CRITICAL: Check if model needs vocabulary expansion after loading new tokenizer
                        if self.model is not None:
                            self.model.tokenizer = self.tokenizer
                            self.model.check_and_expand_vocab_if_needed()
                            # Update layer pool's vocab_capacity for future layer calculations
                            if hasattr(self, 'layer_pool') and self.layer_pool:
                                self.layer_pool.vocab_capacity = self.model.vocab_capacity
                    else:
                        logger.debug(f"[TOKENIZER] Already up to date: {self.tokenizer.current_vocab_size} tokens")
                    return
            except requests.RequestException as e:
                logger.debug(f"[TOKENIZER] CDN fetch failed: {e}")
            
            # Fallback to cached version
            if os.path.exists(tokenizer_cache_path) and self.tokenizer.next_merge_id <= 266:
                try:
                    from neuroshard.core.model.tokenizer import NeuroTokenizer
                    learned_tokenizer = NeuroTokenizer.load(tokenizer_cache_path)
                    
                    if learned_tokenizer.next_merge_id > self.tokenizer.next_merge_id:
                        self.tokenizer.merges = learned_tokenizer.merges
                        self.tokenizer.merge_to_tokens = learned_tokenizer.merge_to_tokens
                        self.tokenizer.next_merge_id = learned_tokenizer.next_merge_id
                        logger.info(f"[TOKENIZER] Loaded cached BPE tokenizer: {self.tokenizer.current_vocab_size} tokens")
                        
                        # CRITICAL: Check if model needs vocabulary expansion
                        if self.model is not None:
                            self.model.tokenizer = self.tokenizer
                            self.model.check_and_expand_vocab_if_needed()
                            # Update layer pool's vocab_capacity for future layer calculations
                            if hasattr(self, 'layer_pool') and self.layer_pool:
                                self.layer_pool.vocab_capacity = self.model.vocab_capacity
                except Exception as e:
                    logger.warning(f"[TOKENIZER] Failed to load cached tokenizer: {e}")
                    
        except Exception as e:
            logger.warning(f"[TOKENIZER] Error loading learned tokenizer: {e}")
    
    def _setup_training(self):
        """Setup training components."""
        from neuroshard.core.training.distributed import FederatedDataManager
        
        # Collect all parameters from my layers
        all_params = []
        for layer in self.model.my_layers.values():
            all_params.extend(layer.parameters())
        if self.model.embedding:
            all_params.extend(self.model.embedding.parameters())
        if self.model.lm_head:
            all_params.extend(self.model.lm_head.parameters())
        if self.model.final_norm:
            all_params.extend(self.model.final_norm.parameters())
        
        self.optimizer = torch.optim.AdamW(all_params, lr=1e-4, weight_decay=0.01)
        
        self.data_manager = FederatedDataManager(
            tokenizer=self.tokenizer,
            max_seq_len=2048
        )
        
        # DYNAMIC TRAINING CONFIG: Calculate based on current model size and device
        # This will be recalculated when model grows via recalculate_training_config()
        num_layers = len(self.my_layer_ids)
        
        # Smart gradient checkpointing decision
        # CRITICAL: Always enable checkpointing for models with many layers!
        # Without checkpointing, attention scores alone need: batch × heads × seq² × layers × 4 bytes
        # For 46 layers: 8 × 16 × 2048² × 46 × 4 = ~92GB (way more than any GPU!)
        
        # SIMPLE RULE: Enable checkpointing if layers > 16 (always safe)
        # This avoids complex calculations that can have bugs with timing of vocab expansion
        if num_layers > 16:
            self._use_gradient_checkpointing = True
            logger.info(f"[NODE] Gradient checkpointing: ENABLED (layers={num_layers} > 16)")
        elif self.device != "cuda":
            # CPU/MPS always use checkpointing for memory efficiency
            self._use_gradient_checkpointing = True
            logger.info(f"[NODE] Gradient checkpointing: ENABLED (device={self.device})")
        else:
            # Small CUDA models can skip checkpointing for speed
            self._use_gradient_checkpointing = False
            logger.info(f"[NODE] Gradient checkpointing: DISABLED (layers={num_layers} ≤ 16, CUDA)")
        
        # Calculate memory-aware training batch size
        self._training_batch_size = self._calculate_training_batch_size()

        logger.info(f"Training initialized: batch_size={self._training_batch_size}, "
                   f"checkpointing={self._use_gradient_checkpointing}, "
                   f"layers={num_layers}, device={self.device}")
        
        # CUDA sanity check: verify GPU is actually usable
        if self.device == "cuda":
            try:
                import time as _time
                test_tensor = torch.randn(1000, 1000, device="cuda")
                start = _time.time()
                _ = torch.matmul(test_tensor, test_tensor)
                torch.cuda.synchronize()
                elapsed = _time.time() - start
                del test_tensor
                torch.cuda.empty_cache()
                logger.info(f"[CUDA] GPU sanity check passed: 1000x1000 matmul in {elapsed*1000:.1f}ms")
            except Exception as e:
                logger.error(f"[CUDA] GPU sanity check FAILED: {e}")
                logger.error("[CUDA] Training will likely run on CPU despite device=cuda!")
    
    def _calculate_training_batch_size(self) -> int:
        """
        Calculate optimal batch size based on available memory, device, and model size.
        
        DYNAMIC: This is called initially and can be recalculated when model grows.
        SMART: Considers GPU memory, gradient checkpointing, and actual model size.
        """
        seq_len = 512  # Typical sequence length
        hidden_dim = self.layer_pool.current_architecture.hidden_dim
        num_layers = len(self.my_layer_ids)
        
        # Calculate model memory footprint (params + gradients + optimizer states)
        model_params = sum(p.numel() for p in self.model.parameters())
        # Model memory: weights (fp32=4 bytes) × 4 (weights + grads + adam_m + adam_v)
        model_memory_mb = (model_params * 4 * 4) / (1024 * 1024)
        
        # For CUDA, check actual GPU memory available
        if self.device == "cuda":
            try:
                gpu_total = torch.cuda.get_device_properties(0).total_memory / (1024 * 1024)
                gpu_allocated = torch.cuda.memory_allocated(0) / (1024 * 1024)
                logger.info(f"[NODE] CUDA memory: {gpu_allocated:.0f}MB used / {gpu_total:.0f}MB total")
                effective_memory_mb = self.available_memory_mb
            except Exception:
                effective_memory_mb = self.available_memory_mb
        else:
            effective_memory_mb = self.available_memory_mb
        
        # CORRECT FORMULA: Available for activations = Total - Model memory
        # Leave 10% buffer for system overhead
        available_for_activations = max(100, (effective_memory_mb * 0.9) - model_memory_mb)
        
        # With gradient checkpointing, activation memory is MUCH lower
        use_checkpointing = getattr(self, '_use_gradient_checkpointing', False)
        if use_checkpointing:
            # Checkpointing: Only need to store ~sqrt(num_layers) worth of activations
            # Plus inputs/outputs at checkpoint boundaries
            checkpoint_segments = max(1, int(num_layers ** 0.5))
            # Memory per sample: seq_len × hidden_dim × checkpoint_segments × 4 bytes × 2 (fwd+bwd)
            mem_per_sample_mb = (seq_len * hidden_dim * checkpoint_segments * 4 * 2) / (1024 * 1024)
            logger.info(f"[NODE] Gradient checkpointing: {checkpoint_segments} segments "
                       f"(~{mem_per_sample_mb:.1f}MB/sample)")
        else:
            # No checkpointing: full activation memory for all layers
            mem_per_sample_mb = (seq_len * hidden_dim * num_layers * 4 * 2) / (1024 * 1024)
        
        logger.info(f"[NODE] Memory budget: total={effective_memory_mb:.0f}MB, "
                   f"model={model_memory_mb:.0f}MB, "
                   f"available_for_activations={available_for_activations:.0f}MB")
        
        # Calculate max batch size from available memory
        max_batch = max(1, int(available_for_activations / max(1, mem_per_sample_mb)))
        
        # SMART CLAMPING based on device capability
        if self.device == "cuda" and effective_memory_mb > 16000:
            # High-memory CUDA (Jetson Orin 32GB, RTX 3090 24GB): up to 8
            max_batch = min(max_batch, 8)
        elif self.device == "cuda" and effective_memory_mb > 8000:
            # Medium CUDA: up to 4
            max_batch = min(max_batch, 4)
        elif self.device == "cuda":
            # Small CUDA: up to 2
            max_batch = min(max_batch, 2)
        elif num_layers > 100:
            # Large model on CPU/MPS: conservative
            max_batch = min(max_batch, 2)
        else:
            max_batch = min(max_batch, 4)
        
        batch_size = max(1, max_batch)
        
        logger.info(f"[NODE] Training config: batch_size={batch_size}, "
                   f"model={model_params/1e6:.1f}M params ({num_layers} layers × {hidden_dim} dim), "
                   f"checkpointing={use_checkpointing}, device={self.device}")
        
        return batch_size
    
    def recalculate_training_config(self):
        """
        Recalculate training configuration after model architecture changes.
        
        Called when:
        - Model grows (new layers added)
        - Memory allocation changes
        - Device changes
        """
        old_batch = getattr(self, '_training_batch_size', None)
        self._training_batch_size = self._calculate_training_batch_size()
        
        # Update gradient checkpointing based on new model size
        num_layers = len(self.my_layer_ids)
        old_checkpointing = getattr(self, '_use_gradient_checkpointing', False)
        
        # Simple checkpointing rule: enable if layers > 16
        if num_layers > 16:
            self._use_gradient_checkpointing = True
        elif self.device != "cuda":
            self._use_gradient_checkpointing = True
        else:
            self._use_gradient_checkpointing = False
        
        if old_batch != self._training_batch_size or old_checkpointing != self._use_gradient_checkpointing:
            logger.info(f"[NODE] Training config updated: batch_size={old_batch}→{self._training_batch_size}, "
                       f"checkpointing={old_checkpointing}→{self._use_gradient_checkpointing}")
    
    def stop(self):
        """Stop the node."""
        logger.info("Stopping DynamicNeuroNode...")
        
        self.is_running = False
        
        # Unregister from network
        if self.layer_pool:
            self.layer_pool.unregister_node(self.node_id)
        
        # Save checkpoint
        self._save_checkpoint()
        
        logger.info("DynamicNeuroNode stopped")
    
    def connect_p2p(self, p2p_manager):
        """Connect to P2P network."""
        self.p2p_manager = p2p_manager

        # Initialize Data Swarm
        from neuroshard.core.network.p2p_data import DataSwarm

        # Ensure cache dir exists in a writable location
        data_cache_dir = self.CHECKPOINT_DIR / "data_cache"
        data_cache_dir.mkdir(parents=True, exist_ok=True)

        self.swarm = DataSwarm(p2p_manager, cache_dir=str(data_cache_dir))

        # Update layer pool with DHT
        if self.layer_pool and hasattr(p2p_manager, 'dht'):
            self.layer_pool.dht = p2p_manager.dht
        
        # IMPORTANT: Give model access to p2p_manager for dynamic layer updates
        # When vocab expansion removes layers, the model needs to update DHT
        if self.model:
            self.model._p2p_manager = p2p_manager

        logger.info("Connected to P2P network and Data Swarm")
    
    # ==================== INFERENCE ====================
    
    def forward(self, input_ids: torch.Tensor, session_id: Optional[str] = None) -> torch.Tensor:
        """
        Forward pass - routes through network if needed.
        
        If this node has all layers: process locally
        If not: forward to nodes with other layers
        """
        # Check if we can do full inference locally
        capacity = self.layer_pool.get_network_capacity()
        
        if len(self.my_layer_ids) == capacity.assigned_layers and self.model.has_embedding and self.model.has_lm_head:
            # We have everything - do local inference
            return self._forward_local(input_ids)
        else:
            # Need to route through network
            return self._forward_distributed(input_ids, session_id)
    
    def _forward_local(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Full local inference (when we have all layers)."""
        with torch.no_grad():
            # Embed
            hidden = self.model.embed(input_ids.to(self.device))
            
            # Forward through all layers
            hidden = self.model.forward_my_layers(hidden)
            
            # Compute logits
            logits = self.model.compute_logits(hidden)
            
            self.inference_count += 1
            self.total_tokens_processed += input_ids.numel()
            
            return logits
    
    def _forward_distributed(self, input_ids: torch.Tensor, session_id: Optional[str] = None) -> torch.Tensor:
        """Distributed inference through network pipeline."""
        # Get pipeline route
        route = self.layer_pool.get_pipeline_route()
        
        if not route:
            raise RuntimeError("No pipeline route available")
        
        # Start with embedding
        if self.model.has_embedding:
            hidden = self.model.embed(input_ids.to(self.device))
        else:
            # Request embedding from holder
            hidden = self._request_embedding(input_ids)
        
        # Forward through layers (local or remote)
        current_layer = 0
        for layer_id, grpc_addr in route:
            if layer_id in self.model.my_layers:
                # Local layer
                hidden, _ = self.model.my_layers[layer_id](hidden)
            else:
                # Remote layer - forward to peer
                hidden = self._forward_to_peer(grpc_addr, hidden, layer_id)
            current_layer = layer_id
        
        # Compute logits
        if self.model.has_lm_head:
            logits = self.model.compute_logits(hidden)
        else:
            # Request from holder
            logits = self._request_logits(hidden)
        
        self.inference_count += 1
        self.total_tokens_processed += input_ids.numel()
        
        return logits
    
    def forward_pipeline(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        training_labels: Optional[torch.Tensor] = None,
        session_id: Optional[str] = None,
        sender_url: Optional[str] = None,
        use_cache: bool = False
    ) -> Tuple[torch.Tensor, Optional[List[torch.Tensor]]]:
        """
        Forward pass for pipeline parallelism (received from peer).
        """
        # Enable gradient tracking if training
        is_training = training_labels is not None
        
        if is_training:
            hidden_states.requires_grad_(True)
            hidden_states.retain_grad()
        
        # Check if input is token IDs (embedding request)
        # Integer dtype or 2D shape [batch, seq] implies input_ids
        # This happens when a client sends input_ids to the Driver (Layer 0)
        if (hidden_states.dtype in [torch.long, torch.int64, torch.int32] or 
            len(hidden_states.shape) == 2) and self.model.has_embedding:
            
            # Ensure correct dtype
            if hidden_states.dtype != torch.long:
                hidden_states = hidden_states.to(torch.long)
            
            # Embed tokens
            hidden_states = self.model.embed(hidden_states)
            
            if is_training:
                hidden_states.requires_grad_(True)
                hidden_states.retain_grad()
        
        # Forward through local layers
        output = self.model.forward_my_layers(hidden_states)
        
        if is_training and session_id:
            # Save context for backward pass
            self.training_context[session_id] = {
                "input": hidden_states,
                "output": output,
                "sender_url": sender_url,
                "timestamp": time.time()
            }
            # Cleanup old sessions
            now = time.time()
            to_remove = [s for s, ctx in self.training_context.items() if now - ctx["timestamp"] > 600]
            for s in to_remove:
                del self.training_context[s]
        
        # If we are the Validator (Last Layer holder)
        # DYNAMIC CHECK: Query layer_pool for current lm_head_holder
        # This handles the case where a new Validator joined and took over
        is_current_validator = self.model.has_lm_head
        if hasattr(self, 'layer_pool') and self.layer_pool:
            is_current_validator = (self.layer_pool.lm_head_holder == self.node_id)
        
        if is_current_validator:
            logits = self.model.compute_logits(output)
            
            # Calculate Loss if labels present
            if training_labels is not None:
                loss = torch.nn.functional.cross_entropy(
                    logits.view(-1, logits.size(-1)),
                    training_labels.view(-1),
                    ignore_index=-100
                )
                
                # Use training lock to prevent conflict with local training
                with self._training_lock:
                    # Trigger Backward Pass
                    self.optimizer.zero_grad()
                    loss.backward()
                    
                    # Propagate gradient back to previous node
                    if sender_url and session_id:
                        # The gradient we send back is dL/d(input_hidden_states)
                        # hidden_states.grad is populated by backward()
                        if hidden_states.grad is not None:
                            self._backward_to_peer(
                                sender_url, 
                                hidden_states.grad, 
                                # Target shard is whatever layer sent this to us. 
                                # Assuming sender holds previous layers.
                                # We send to the sender's LAST layer.
                                # Simplified: just send to the node, it routes.
                                0, 
                                session_id
                            )
                    
                    # Step Optimizer
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.optimizer.step()
                
                self.total_training_rounds += 1
                self.current_loss = loss.item()
                
                return logits, None
            
            return logits, None
            
        # If we are a Worker (Middle Layer), we need to forward to next peer
        my_last_layer = max(self.my_layer_ids) if self.my_layer_ids else 0
        next_layer = my_last_layer + 1
        
        if self.p2p_manager:
            next_hop = self.p2p_manager.get_next_hop(next_layer)
            if next_hop:
                return self._forward_to_peer(
                    next_hop, 
                    output, 
                    next_layer, 
                    labels=training_labels, 
                    session_id=session_id
                )
        
        logger.warning(f"Pipeline broken at layer {next_layer}: no peer found")
        return output, None

    def backward_pipeline(self, grad_output: torch.Tensor, session_id: str):
        """
        Backward pass received from next peer.
        """
        if session_id not in self.training_context:
            logger.warning(f"Received backward for unknown session {session_id}")
            return
            
        ctx = self.training_context[session_id]
        output = ctx["output"]
        input_tensor = ctx["input"]
        sender_url = ctx["sender_url"]
        
        # Use training lock to prevent conflict with local training
        with self._training_lock:
            # Run local backward
            # output is the tensor we produced in forward_pipeline
            # grad_output is dL/d(output) received from next peer
            self.optimizer.zero_grad()
            output.backward(grad_output)
            
            # Propagate back
            if sender_url and input_tensor.grad is not None:
                 # Find previous layer ID? Not strictly needed for routing if we have direct sender URL
                 # But _backward_to_peer takes layer_id
                 self._backward_to_peer(sender_url, input_tensor.grad, 0, session_id)
                 
            # Step Optimizer
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()
        
        # Cleanup
        del self.training_context[session_id]

    def _request_embedding(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Request embedding from the node that holds it."""
        # Find a node that holds Layer 0 (Driver)
        peer_url = None
        
        # 1. Check layer pool assignments
        if self.layer_pool:
            assignments = self.layer_pool.get_layer_holders(0)
            if assignments:
                # Pick one (e.g., random for load balancing)
                import random
                peer_url = random.choice(assignments).grpc_addr
        
        # 2. Fallback to P2P manager routing
        if not peer_url and self.p2p_manager:
            peer_url = self.p2p_manager.get_next_hop(0)
            
        if not peer_url:
            raise RuntimeError("No embedding holder (Driver/Layer 0) found in network")
            
        # Call peer - Send input_ids to Layer 0 holder
        # The receiver's forward_pipeline will detect it's input_ids and run embed()
        result, _ = self._forward_to_peer(peer_url, input_ids, 0)
        return result
    
    def _forward_to_peer(self, peer_url: str, hidden: torch.Tensor, layer_id: int, labels: Optional[torch.Tensor] = None, session_id: str = None) -> torch.Tensor:
        """
        Forward hidden states to a peer for processing.
        
        SECURITY: Calculates and validates SHA256 checksums to detect tampering.
        """
        from protos import neuroshard_pb2
        from protos import neuroshard_pb2_grpc
        from neuroshard.core.network.connection_pool import get_channel
        import numpy as np
        import hashlib
        
        try:
            parsed = urlparse(peer_url)
            ip = parsed.hostname
            # gRPC port convention
            port = (parsed.port or 80) + 1000
            
            channel = get_channel(f"{ip}:{port}")
            stub = neuroshard_pb2_grpc.NeuroShardServiceStub(channel)
            
            # Serialize hidden states
            hidden_bytes = hidden.detach().cpu().numpy().tobytes()
            hidden_shape = list(hidden.shape)
            
            # CHECKSUM: Calculate SHA256 hash for integrity verification
            checksum = hashlib.sha256(hidden_bytes).hexdigest()
            logger.debug(f"[SECURITY] Sending layer {layer_id} with checksum: {checksum[:16]}...")
            
            # Serialize labels if present
            labels_bytes = b""
            if labels is not None:
                labels_bytes = labels.cpu().numpy().tobytes()
            
            req_session_id = session_id or f"train_{time.time()}"
            
            # Get my URL for backward routing
            my_url = ""
            if self.p2p_manager:
                my_url = self.p2p_manager.my_url
            
            req = neuroshard_pb2.PipelineForwardRequest(
                session_id=req_session_id,
                request_id=f"req_{time.time()}",
                hidden_states=hidden_bytes,
                hidden_shape=hidden_shape,
                target_shard=layer_id,
                use_cache=False,
                training_labels=labels_bytes,
                sender_url=my_url
            )
            
            # Store context for backward pass
            # We need to know WHO sent us this so we can send gradients back? 
            # No, this function is called by US sending to THEM.
            # We need to know who THEY are so when they send us gradients back, we verify?
            # Actually, we don't need to do anything here for backward. 
            # They will call PipelineBackward on US.
            
            resp = stub.PipelineForward(req, timeout=30.0)
            
            if not resp.success:
                raise RuntimeError(f"Peer error: {resp.error_message}")
            
            # Deserialize result
            if resp.is_final:
                # It's logits
                result_bytes = resp.logits
                result = torch.from_numpy(
                    np.frombuffer(result_bytes, dtype=np.float32)
                ).reshape(list(resp.logits_shape))
            else:
                # It's hidden states (recursive/chained)
                result_bytes = resp.hidden_states
                result = torch.from_numpy(
                    np.frombuffer(result_bytes, dtype=np.float32)
                ).reshape(list(resp.hidden_shape))
            
            # CHECKSUM VALIDATION: Verify integrity of received data
            received_checksum = hashlib.sha256(result_bytes).hexdigest()
            logger.debug(f"[SECURITY] Received layer {layer_id} result with checksum: {received_checksum[:16]}...")
            
            # AUDIT TRAIL: Store checksum in PipelineSession for tamper detection
            if session_id and self.ledger and hasattr(self.ledger, 'inference_market'):
                market = self.ledger.inference_market
                if market and hasattr(market, 'active_sessions'):
                    for sess_id, session in market.active_sessions.items():
                        if sess_id == session_id or session.request_id in session_id:
                            session.activations_hashes.append(received_checksum)
                            logger.debug(f"[AUDIT] Stored checksum for layer {layer_id} in session")
                            break
            
            return result.to(self.device), None
            
        except Exception as e:
            logger.error(f"Failed to forward to peer {peer_url}: {e}")
            return hidden, None

    def _backward_to_peer(self, peer_url: str, grad_output: torch.Tensor, layer_id: int, session_id: str):
        """Send gradients back to the previous peer."""
        from protos import neuroshard_pb2
        from protos import neuroshard_pb2_grpc
        from neuroshard.core.network.connection_pool import get_channel
        
        try:
            parsed = urlparse(peer_url)
            ip = parsed.hostname
            port = (parsed.port or 80) + 1000
            
            channel = get_channel(f"{ip}:{port}")
            stub = neuroshard_pb2_grpc.NeuroShardServiceStub(channel)
            
            grad_bytes = grad_output.detach().cpu().numpy().tobytes()
            grad_shape = list(grad_output.shape)
            
            req = neuroshard_pb2.PipelineBackwardRequest(
                session_id=session_id,
                request_id=f"bw_{time.time()}",
                grad_output=grad_bytes,
                grad_shape=grad_shape,
                target_shard=layer_id
            )
            
            stub.PipelineBackward(req, timeout=10.0)
            
        except Exception as e:
            logger.error(f"Failed to backward to peer {peer_url}: {e}")

    def _request_logits(self, hidden: torch.Tensor) -> torch.Tensor:
        """Request logits from the node that holds LM head."""
        # Find Last Layer holder (Validator)
        if not self.layer_pool:
             return hidden
             
        capacity = self.layer_pool.get_network_capacity()
        last_layer = max(0, capacity.assigned_layers - 1)
        
        peer_url = None
        
        # 1. Check layer pool assignments
        assignments = self.layer_pool.get_layer_holders(last_layer)
        if assignments:
            import random
            peer_url = random.choice(assignments).grpc_addr
            
        # 2. Fallback to P2P manager
        if not peer_url and self.p2p_manager:
            peer_url = self.p2p_manager.get_next_hop(last_layer)
            
        if not peer_url:
            raise RuntimeError(f"No Validator (Layer {last_layer}) found in network")
            
        # Forward hidden states to peer targeting Last Layer
        # The receiver will compute logits and return is_final=True
        return self._forward_to_peer(peer_url, hidden, last_layer)
    
    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 50,
        temperature: float = 1.0,
    ) -> str:
        """Generate text from prompt."""
        try:
            if not self.tokenizer:
                raise RuntimeError("Tokenizer not initialized")
            
            input_ids = torch.tensor([self.tokenizer.encode(prompt)], dtype=torch.long)
            logger.debug(f"[GENERATE] Encoded prompt: {input_ids.shape} tokens")
            
            # Move to model's device (handles CPU, CUDA, MPS)
            generated = input_ids.clone().to(self.device)
            
            # Get current vocabulary size from tokenizer
            # Only tokens 0 to current_vocab_size-1 are valid (have learned representations)
            # This is NOT a workaround - it's how BPE tokenizers work (vocab grows over time)
            valid_vocab_size = self.tokenizer.current_vocab_size
            
            for step in range(max_new_tokens):
                logits = self.forward(generated)
                next_logits = logits[:, -1, :] / temperature
                
                # Constrain to valid vocabulary (standard BPE tokenizer behavior)
                # Tokens beyond current_vocab_size don't exist in the tokenizer yet
                if valid_vocab_size < next_logits.size(-1):
                    next_logits[:, valid_vocab_size:] = float('-inf')
                
                probs = torch.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
                generated = torch.cat([generated, next_token], dim=-1)
                
                if next_token.item() == 2:  # EOS
                    logger.debug(f"[GENERATE] EOS at step {step+1}")
                    break
            
            prompt_tokens = input_ids.size(1)
            new_tokens = generated[0, prompt_tokens:].tolist()
            result = self.tokenizer.decode(new_tokens)
            logger.debug(f"[GENERATE] Generated {len(new_tokens)} tokens: '{result[:100]}...'")
            
            return result
            
        except Exception as e:
            logger.error(f"[GENERATE] Error: {e}")
            import traceback
            logger.error(traceback.format_exc())
            raise
    
    # ==================== TRAINING ====================
    
    def contribute_training_data(self, text: str, apply_dp: bool = True) -> int:
        """
        Contribute training data.
        
        Returns the number of tokens added.
        """
        if not self.data_manager:
            return 0
        
        # Get token count before
        stats_before = self.data_manager.get_stats()
        tokens_before = stats_before.get("total_tokens", 0)
        
        self.data_manager.add_text(text, apply_dp=apply_dp)
        
        # Get token count after
        stats_after = self.data_manager.get_stats()
        tokens_after = stats_after.get("total_tokens", 0)
        
        tokens_added = tokens_after - tokens_before
        logger.info(f"Added {tokens_added} tokens to training buffer")
        
        return tokens_added
    
    def _get_training_batch(self) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
        """
        Get a training batch from the Genesis data loader.
        
        Returns:
            Tuple of (input_ids, labels) or None if data not available.
            
        Note: Only Drivers (with embedding) load training data.
              Workers wait for activations via pipeline, they don't need data directly.
        """
        if not self.enable_training:
            return None
        
        # WORKERS DON'T LOAD DATA - they receive activations via pipeline
        # Only Drivers (with embedding) need to load training data
        if not self.model.has_embedding:
            return None  # Worker - skip training data loading
        
        # Initialize genesis loader if needed
        if not hasattr(self, 'genesis_loader') or self.genesis_loader is None:
            try:
                from neuroshard.core.training.distributed import GenesisDataLoader
                from neuroshard.core.model.tokenizer import get_neuro_tokenizer
                logger.info("[GENESIS] Initializing data loader...")
                self.genesis_loader = GenesisDataLoader(
                    self.node_id, 
                    get_neuro_tokenizer(),
                    max_storage_mb=self.max_storage_mb
                )
                logger.info(f"[GENESIS] Data loader ready: {self.genesis_loader.total_shards} shards available")
                
                # Connect Swarm to Loader
                if hasattr(self, 'swarm') and self.swarm:
                    self.genesis_loader.set_swarm(self.swarm)
            except Exception as e:
                logger.warning(f"[GENESIS] Failed to initialize loader: {e}")
                return None
        
        # Check if data is ready
        if not self.genesis_loader.is_data_ready():
            return None
        
        # Get batch
        batch_size = getattr(self, '_training_batch_size', 2)
        try:
            input_ids, labels = self.genesis_loader.get_batch(batch_size=batch_size)
            return input_ids, labels
        except Exception as e:
            logger.warning(f"[GENESIS] Failed to get batch: {e}")
            return None
    
    def train_step(self) -> Optional[float]:
        """
        Perform a training step on my layers.
        
        OPTIMIZED FOR SINGLE-NODE: When we have embedding + all layers + LM head,
        we skip distributed overhead and train locally.
        
        NON-BLOCKING DATA: Uses prefetched data when available, raises RuntimeError
        if data not ready (caller should retry later).
        """
        if not self.enable_training:
            return None
        
        # RUNTIME MEMORY CHECK: Skip training if system memory is critically high
        # This prevents OOM crashes and keeps the system responsive
        try:
            import psutil
            mem = psutil.virtual_memory()
            # Skip if less than 15% of system RAM is free (critical threshold)
            if mem.percent > 85:
                logger.warning(f"[NODE] System memory at {mem.percent:.0f}%, skipping training step")
                # Also try to free some memory
                import gc
                gc.collect()
                if self.device == "cuda":
                    torch.cuda.empty_cache()
                elif self.device == "mps":
                    torch.mps.empty_cache()
                return None
        except Exception:
            pass  # If psutil fails, continue anyway
        
        try:
            # SINGLE-NODE OPTIMIZATION: Check if we're a full node (Driver + Worker + Validator)
            # DYNAMIC CHECK: Use layer_pool to get current lm_head_holder
            # This handles the case where a new Validator joined and took over the LM head
            am_current_validator = self.model.has_lm_head
            if hasattr(self, 'layer_pool') and self.layer_pool:
                am_current_validator = (self.layer_pool.lm_head_holder == self.node_id)
            
            is_full_node = self.model.has_embedding and am_current_validator
            
            if self.model.has_embedding:
                # I am a Driver (Layer 0)
                # Use Genesis Data Loader
                if not hasattr(self, 'genesis_loader') or self.genesis_loader is None:
                    try:
                        from neuroshard.core.training.distributed import GenesisDataLoader
                        from neuroshard.core.model.tokenizer import get_neuro_tokenizer
                        logger.info("[GENESIS] Initializing data loader...")
                        self.genesis_loader = GenesisDataLoader(
                            self.node_id, 
                            get_neuro_tokenizer(),
                            max_storage_mb=self.max_storage_mb
                        )
                        logger.info(f"[GENESIS] Data loader ready: {self.genesis_loader.total_shards} shards available")
                        
                        # Connect Swarm to Loader
                        if hasattr(self, 'swarm') and self.swarm:
                            self.genesis_loader.set_swarm(self.swarm)
                    except Exception as e:
                        import traceback
                        logger.error(f"[GENESIS] ERROR: {type(e).__name__}: {e}")
                        logger.error(f"[GENESIS] {traceback.format_exc()}")
                        # Mark as failed so we don't keep retrying immediately
                        self.genesis_loader = None
                        raise RuntimeError(f"Genesis loader init failed: {e}")
                
                # Check if data is ready (non-blocking)
                if not self.genesis_loader.is_data_ready():
                    # Data not ready - don't block, let caller retry
                    raise RuntimeError("Data not ready - shard still loading")
                
                # Get batch from Genesis Shard using memory-aware batch size
                batch_size = getattr(self, '_training_batch_size', 2)
                try:
                    input_ids, labels = self.genesis_loader.get_batch(batch_size=batch_size)
                    input_ids = input_ids.to(self.device)
                    labels = labels.to(self.device)
                except RuntimeError as e:
                    # Data not ready - propagate up
                    raise
                except Exception as e:
                    logger.warning(f"[GENESIS] Failed to get batch: {type(e).__name__}: {e}")
                    import traceback
                    logger.warning(traceback.format_exc())
                    return None
                
                # SINGLE-NODE OPTIMIZED PATH: Skip distributed overhead
                if is_full_node:
                    return self._train_step_local(input_ids, labels)
                
                # DISTRIBUTED PATH: Forward to next peer
                # Forward pass with optional gradient checkpointing
                # Note: time.sleep(0) yields GIL to keep HTTP server responsive
                embeddings = self.model.embed(input_ids)
                embeddings.requires_grad_(True)
                embeddings.retain_grad()
                time.sleep(0)  # Yield GIL
                
                # Use gradient checkpointing if enabled (trades CPU for memory)
                if getattr(self, '_use_gradient_checkpointing', False):
                    output = torch.utils.checkpoint.checkpoint(
                        self.model.forward_my_layers,
                        embeddings,
                        use_reentrant=False
                    )
                else:
                    output = self.model.forward_my_layers(embeddings)
                time.sleep(0)  # Yield GIL after forward pass
                
                # Distributed: Send to next peer
                my_last_layer = max(self.my_layer_ids) if self.my_layer_ids else 0
                next_layer = my_last_layer + 1
                
                if self.p2p_manager:
                    next_hop = self.p2p_manager.get_next_hop(next_layer)
                    if next_hop:
                        session_id = f"train_{self.node_id}_{time.time()}"
                        
                        # Save context for backward
                        self.training_context[session_id] = {
                            "input": embeddings,
                            "output": output,
                            "sender_url": None, # We are the start
                            "timestamp": time.time()
                        }
                        
                        result, _ = self._forward_to_peer(
                            next_hop, 
                            output, 
                            next_layer, 
                            labels=labels, 
                            session_id=session_id
                        )
                        
                        # Check if forward succeeded (result should be different from output if it was processed)
                        # If the peer rejected or failed, result will be the original output (unchanged)
                        forward_succeeded = result is not output
                        
                        if not forward_succeeded:
                            # Pipeline forward failed - peer rejected or error
                            # Clean up the training context
                            if session_id in self.training_context:
                                del self.training_context[session_id]
                            logger.warning(f"[DISTRIBUTED] Pipeline forward failed - skipping training step")
                            return None
                        
                        # We don't get loss immediately in distributed pipeline
                        # It comes back later via backward pass or status update
                        # For now, return None (not inf!)
                        return None
                
                return None
                
            else:
                # I am a Worker/Validator
                # I wait for activations from peers via gRPC (forward_pipeline)
                # So this method does nothing actively
                return None
                
        except RuntimeError as e:
            error_msg = str(e)
            if "not ready" in error_msg.lower():
                # Data not ready - propagate to caller
                raise
            elif "out of memory" in error_msg.lower() or "MPS" in error_msg:
                logger.warning(f"Training step OOM - reducing batch size and clearing cache")
                # Clear GPU cache
                import gc
                gc.collect()
                if self.device == "mps":
                    torch.mps.empty_cache()
                elif self.device == "cuda":
                    torch.cuda.empty_cache()
                
                # Reduce batch size for next attempt
                current_batch = getattr(self, '_training_batch_size', 8)
                if current_batch > 1:
                    self._training_batch_size = max(1, current_batch // 2)
                    logger.info(f"Reduced batch size to {self._training_batch_size}")
                else:
                    # Already at minimum batch size, fall back to CPU for training
                    if self.device != "cpu":
                        logger.warning(f"Batch size already at minimum. Consider using --memory flag to limit layers.")
            else:
                logger.error(f"Training step failed: {e}")
            return None
        except Exception as e:
            logger.error(f"Training step failed: {e}")
            return None
    
    def _train_step_local(self, input_ids: torch.Tensor, labels: torch.Tensor) -> float:
        """
        OPTIMIZED single-node training step.
        
        When we have ALL components (embedding + layers + LM head), we can
        train entirely locally without any network overhead.
        """
        # DIAGNOSTIC: Verify device placement periodically
        if self.total_training_rounds % 100 == 0:
            try:
                emb_device = next(self.model.embedding.parameters()).device if self.model.embedding else 'N/A'
                layer_device = next(iter(self.model.my_layers.values())).parameters().__next__().device if self.model.my_layers else 'N/A'
                logger.info(f"[TRAIN] Device check: input={input_ids.device}, embedding={emb_device}, layer0={layer_device}")
            except Exception as e:
                logger.warning(f"[TRAIN] Device check failed: {e}")
        
        # Forward pass with optional gradient checkpointing
        embeddings = self.model.embed(input_ids)
        
        # Use gradient checkpointing if enabled (trades CPU for memory)
        if getattr(self, '_use_gradient_checkpointing', False):
            output = torch.utils.checkpoint.checkpoint(
                self.model.forward_my_layers,
                embeddings,
                use_reentrant=False
            )
        else:
            output = self.model.forward_my_layers(embeddings)
        
        # Compute logits and loss
        logits = self.model.compute_logits(output)
        
        loss = torch.nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)),
            labels.view(-1),
            ignore_index=-100
        )
        
        # Use training lock to prevent conflict with pipeline training
        with self._training_lock:
            # Backward pass
            self.optimizer.zero_grad()
            loss.backward()
            
            # Gradient clipping and optimizer step
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()
        
        # Update stats
        self.total_training_rounds += 1
        loss_val = loss.item()
        self.current_loss = loss_val
        
        # PERIODIC CHECKPOINT: Save every 100 steps
        # Synchronous save blocks training briefly (~30-60s) but avoids memory pressure
        if self.total_training_rounds % 100 == 0:
            self._save_checkpoint()
        
        return loss_val
    
    # ==================== STATS & PONW ====================
    
    def get_stats(self) -> Dict[str, Any]:
        """Get node statistics."""
        # Safety check for shutdown race condition
        model = getattr(self, 'model', None)
        layer_pool = getattr(self, 'layer_pool', None)
        
        contribution = model.get_my_contribution() if model else {}
        capacity = layer_pool.get_network_capacity() if layer_pool else None
        
        # Calculate reward multiplier
        my_layer_ids = getattr(self, 'my_layer_ids', [])
        network_layers = capacity.assigned_layers if capacity else len(my_layer_ids)
        reward_multiplier = calculate_reward_multiplier(
            num_layers_held=len(my_layer_ids),
            total_network_layers=network_layers or 1,
            has_embedding=model.has_embedding if model else False,
            has_lm_head=model.has_lm_head if model else False,
        )
        
        # Estimate network params (rough: ~10M params per layer)
        network_params = network_layers * 10_000_000 if network_layers else 0
        
        # Get data buffer size
        data_buffer_size = 0
        if self.data_manager:
            data_stats = self.data_manager.get_stats()
            data_buffer_size = data_stats.get("buffer_size", 0)
        
        # Get shard stats (if we have a genesis loader)
        shard_stats = {}
        if hasattr(self, 'genesis_loader') and self.genesis_loader:
            shard_stats = self.genesis_loader.get_stats()
        
        # Multi-node identity info
        instance_id = getattr(self, 'instance_id', None)
        wallet_id = getattr(self, 'wallet_id', None)
        
        return {
            "node_id": self.node_id[:16] + "...",
            "instance_id": instance_id,  # Unique per machine+port
            "wallet_id": wallet_id,       # Same across instances with same token
            "available_memory_mb": self.available_memory_mb,
            "my_layers": self.my_layer_ids,
            "my_params": contribution.get("my_params", 0),
            "has_embedding": contribution.get("has_embedding", False),
            "has_lm_head": contribution.get("has_lm_head", False),
            "contribution_ratio": contribution.get("contribution_ratio", 0),
            "reward_multiplier": reward_multiplier,
            "network_layers": network_layers,
            "network_params": network_params,
            "network_nodes": capacity.total_nodes if capacity else 1,
            "total_tokens_processed": self.total_tokens_processed,
            "total_training_rounds": self.total_training_rounds,
            "current_loss": self.current_loss,
            "inference_count": self.inference_count,
            "data_buffer_size": data_buffer_size,
            "shard_stats": shard_stats,
        }
    
    def get_ponw_proof(self) -> Dict[str, Any]:
        """
        Generate Proof of Neural Work.
        
        This proof demonstrates verifiable neural network computation
        and is used for NEURO token rewards.
        """
        contribution = self.model.get_my_contribution() if self.model else {}
        capacity = self.layer_pool.get_network_capacity() if self.layer_pool else None
        
        # Calculate reward multiplier
        multiplier = calculate_reward_multiplier(
            num_layers_held=len(self.my_layer_ids),
            total_network_layers=capacity.assigned_layers if capacity else 1,
            has_embedding=self.model.has_embedding if self.model else False,
            has_lm_head=self.model.has_lm_head if self.model else False,
        )
        
        timestamp = time.time()
        
        # Determine role
        role = "Worker"
        if self.model and self.model.has_embedding:
            role = "Driver"
        elif self.model and self.model.has_lm_head:
            role = "Validator"
        
        proof_data = {
            "node_id": self.node_id,
            "timestamp": timestamp,
            "tokens_processed": self.total_tokens_processed,
            "training_rounds": self.total_training_rounds,
            "training_contributions": self.training_contribution_count,
            "inference_count": self.inference_count,
            "layers_held": len(self.my_layer_ids),
            "layer_ids": self.my_layer_ids,
            "has_embedding": self.model.has_embedding if self.model else False,
            "has_lm_head": self.model.has_lm_head if self.model else False,
            "role": role,
            "reward_multiplier": multiplier,
            "available_memory_mb": self.available_memory_mb,
        }
        
        # Add model hash for verification
        # Use architecture-based hash (consistent with SwarmEnabledDynamicNode._get_model_hash)
        if self.model:
            hasher = hashlib.sha256()
            arch_str = f"{self.model.hidden_dim}:{len(self.my_layer_ids)}:{getattr(self.model, 'num_heads', 0)}"
            hasher.update(arch_str.encode())
            for name, param in sorted(self.model.named_parameters()):
                hasher.update(f"{name}:{list(param.shape)}".encode())
            proof_data["model_hash"] = hasher.hexdigest()[:16]
        
        # Sign the proof
        proof_string = f"{self.node_id}:{timestamp}:{self.total_tokens_processed}:{len(self.my_layer_ids)}:{self.total_training_rounds}"
        if self.node_token:
            # Use HMAC for proper signing
            import hmac
            signature = hmac.new(
                self.node_token.encode(),
                proof_string.encode(),
                hashlib.sha256
            ).hexdigest()
        else:
            signature = hashlib.sha256(proof_string.encode()).hexdigest()
        
        proof_data["signature"] = signature
        
        return proof_data
    
    def _prefetch_vocab_capacity(self):
        """
        Pre-fetch tokenizer vocab size to know how much memory embeddings will need.
        
        This MUST be called before register_node() to ensure accurate layer assignment.
        Without this, we'd assign layers assuming 32K vocab, then OOM when vocab expands to 288K+.
        """
        import requests
        import os
        
        GENESIS_CDN_URL = "https://dwquwt9gkkeil.cloudfront.net"
        cache_dir = os.path.join(os.path.expanduser("~"), ".neuroshard", "data_cache")
        tokenizer_cache_path = os.path.join(cache_dir, "tokenizer.json")
        
        vocab_size = INITIAL_VOCAB_SIZE  # Default fallback
        
        try:
            # Try to fetch vocab size from CDN
            tokenizer_url = f"{GENESIS_CDN_URL}/tokenizer.json"
            resp = requests.get(tokenizer_url, timeout=10)
            
            if resp.status_code == 200:
                data = resp.json()
                vocab_size = data.get("next_merge_id", INITIAL_VOCAB_SIZE)
                
                # Cache locally for faster startup next time
                os.makedirs(cache_dir, exist_ok=True)
                with open(tokenizer_cache_path, 'w') as f:
                    f.write(resp.text)
                
                logger.info(f"[VOCAB] Pre-fetched tokenizer: {vocab_size:,} tokens (for memory calculation)")
            else:
                # Try cached version
                if os.path.exists(tokenizer_cache_path):
                    import json
                    with open(tokenizer_cache_path, 'r') as f:
                        data = json.load(f)
                    vocab_size = data.get("next_merge_id", INITIAL_VOCAB_SIZE)
                    logger.info(f"[VOCAB] Using cached tokenizer: {vocab_size:,} tokens")
        except Exception as e:
            # Try cached version as fallback
            try:
                if os.path.exists(tokenizer_cache_path):
                    import json
                    with open(tokenizer_cache_path, 'r') as f:
                        data = json.load(f)
                    vocab_size = data.get("next_merge_id", INITIAL_VOCAB_SIZE)
                    logger.info(f"[VOCAB] Using cached tokenizer: {vocab_size:,} tokens (CDN unavailable)")
            except Exception:
                pass
            logger.debug(f"[VOCAB] Could not prefetch vocab size: {e}, using default {INITIAL_VOCAB_SIZE}")
        
        # Round up to next chunk boundary (no headroom - recalculate if vocab grows)
        # Previously used 64K headroom but this wastes ~1GB memory on limited devices
        vocab_capacity = ((vocab_size + VOCAB_GROWTH_CHUNK - 1) // VOCAB_GROWTH_CHUNK) * VOCAB_GROWTH_CHUNK
        
        # Update layer pool's vocab_capacity for accurate layer assignment
        self.layer_pool.vocab_capacity = vocab_capacity
        logger.info(f"[VOCAB] Layer pool vocab_capacity set to {vocab_capacity:,} (current vocab: {vocab_size:,})")
    
    def _reconcile_architecture(self):
        """
        Smart architecture reconciliation for rejoining the network.
        
        Handles all scenarios:
        1. Quick restart (same architecture) → Use checkpoint
        2. Network upgraded (larger arch) → Start fresh with network arch
        3. Network downgraded (smaller arch) → Start fresh with network arch
        4. Solo bootstrap (no peers) → Use checkpoint or calculate
        5. First time (no checkpoint) → Query network or calculate
        
        Priority:
        1. Network consensus (if peers available)
        2. Saved checkpoint (if compatible)
        3. Fresh calculation (fallback)
        """
        saved_arch = self._peek_checkpoint_architecture()
        network_arch = self._query_network_architecture()
        
        # Log what we found
        if saved_arch:
            logger.info(f"Saved checkpoint: {saved_arch.num_layers}L × {saved_arch.hidden_dim}H")
        else:
            logger.info(f"No saved checkpoint found")
        
        if network_arch:
            logger.info(f"Network architecture: {network_arch.num_layers}L × {network_arch.hidden_dim}H")
        else:
            logger.info(f"No peers found (solo mode or bootstrap)")
        
        # Decision matrix
        if network_arch and saved_arch:
            # Both exist - compare them
            if self._architectures_compatible(saved_arch, network_arch):
                # Perfect - checkpoint matches network
                logger.info(f"✅ Checkpoint compatible with network - will load checkpoint")
                self.layer_pool.current_architecture = network_arch
                self.layer_pool.current_num_layers = network_arch.num_layers
            else:
                # Mismatch - network takes priority
                logger.warning(f"⚠️ Architecture mismatch!")
                logger.warning(f"   Checkpoint: {saved_arch.num_layers}L × {saved_arch.hidden_dim}H")
                logger.warning(f"   Network: {network_arch.num_layers}L × {network_arch.hidden_dim}H")
                
                # Check if network arch fits in our memory
                network_memory = network_arch.estimate_memory_mb()
                if network_memory <= self.available_memory_mb:
                    logger.warning(f"   → Using NETWORK architecture (checkpoint will be incompatible)")
                    logger.warning(f"   → Your training progress will be preserved but weights reset")
                    self.layer_pool.current_architecture = network_arch
                    self.layer_pool.current_num_layers = network_arch.num_layers
                    # Rename old checkpoint instead of deleting
                    self._archive_incompatible_checkpoint()
                else:
                    logger.error(f"   → Network arch needs {network_memory}MB but you only have {self.available_memory_mb}MB!")
                    logger.error(f"   → This node cannot participate in current network")
                    logger.error(f"   → Falling back to solo mode with checkpoint architecture")
                    self.layer_pool.current_architecture = saved_arch
                    self.layer_pool.current_num_layers = saved_arch.num_layers
        
        elif network_arch:
            # Network exists but no checkpoint - join the network
            network_memory = network_arch.estimate_memory_mb()
            if network_memory <= self.available_memory_mb:
                logger.info(f"✅ Joining network with architecture: {network_arch.num_layers}L × {network_arch.hidden_dim}H")
                self.layer_pool.current_architecture = network_arch
                self.layer_pool.current_num_layers = network_arch.num_layers
            else:
                logger.warning(f"⚠️ Network arch needs {network_memory}MB but you only have {self.available_memory_mb}MB")
                logger.warning(f"   → Will calculate a smaller architecture (may train in isolation)")
                # Let register_node calculate appropriate architecture
        
        elif saved_arch:
            # Checkpoint exists but no network peers (solo mode)
            # IMPORTANT: Check ACTUAL layers in checkpoint, not just architecture's num_layers
            actual_saved_layers = self._get_checkpoint_layer_count()
            if actual_saved_layers and actual_saved_layers > saved_arch.num_layers:
                # Model grew beyond base architecture - calculate memory for actual layers
                memory_per_layer = estimate_memory_per_layer(saved_arch)
                saved_memory = memory_per_layer * actual_saved_layers * 1.1  # 10% overhead
                logger.info(f"Checkpoint has {actual_saved_layers} layers (grew from base {saved_arch.num_layers})")
            else:
                saved_memory = saved_arch.estimate_memory_mb()
                actual_saved_layers = saved_arch.num_layers
            
            if saved_memory <= self.available_memory_mb:
                logger.info(f"✅ Solo mode - using saved checkpoint: {actual_saved_layers}L × {saved_arch.hidden_dim}H")
                self.layer_pool.current_architecture = saved_arch
                self.layer_pool.current_num_layers = actual_saved_layers
            else:
                # Memory too small for saved checkpoint - calculate what we CAN fit
                # Use current vocab_capacity for accurate memory estimation
                vocab_cap = getattr(self.layer_pool, 'vocab_capacity', INITIAL_VOCAB_SIZE)
                max_layers = calculate_layer_assignment(
                    self.available_memory_mb, saved_arch, 
                    safety_factor=0.6, vocab_capacity=vocab_cap,
                    training_mode=True  # Conservative for training
                )
                logger.warning(f"⚠️ Saved checkpoint has {actual_saved_layers} layers ({saved_memory:.0f}MB) "
                              f"but you only have {self.available_memory_mb:.0f}MB")
                logger.warning(f"   → Will use {max_layers} layers (reduced from checkpoint)")
                self.layer_pool.current_architecture = saved_arch
                self.layer_pool.current_num_layers = max_layers
        
        else:
            # No checkpoint, no network - fresh start
            logger.info(f"Fresh start - architecture will be calculated from available memory")
    
    def _query_network_architecture(self) -> Optional[ModelArchitecture]:
        """
        Query the network for the current architecture.
        
        Tries multiple sources:
        1. DHT lookup for architecture announcements
        2. Tracker API for network stats
        3. Direct peer query
        
        Returns None if no peers available (solo mode).
        """
        import requests
        
        # Method 1: Try tracker API first (fastest, most reliable)
        if self.tracker_url:
            try:
                # Query tracker for network architecture
                response = requests.get(
                    f"{self.tracker_url}/network_architecture",
                    timeout=5
                )
                if response.ok:
                    data = response.json()
                    if data.get("hidden_dim"):
                        arch = ModelArchitecture(
                            hidden_dim=data["hidden_dim"],
                            intermediate_dim=data.get("intermediate_dim", int(data["hidden_dim"] * 8 / 3)),
                            num_layers=data.get("num_layers", 12),
                            num_heads=data.get("num_heads", 12),
                            num_kv_heads=data.get("num_kv_heads", 4),
                        )
                        logger.debug(f"Got network architecture from tracker: {arch.num_layers}L × {arch.hidden_dim}H")
                        return arch
            except Exception as e:
                logger.debug(f"Tracker architecture query failed: {e}")
        
        # Method 2: Query known peers directly
        if self.p2p_manager and self.p2p_manager.known_peers:
            for peer_url in list(self.p2p_manager.known_peers.keys())[:3]:
                try:
                    response = requests.get(
                        f"{peer_url}/api/node/architecture",
                        timeout=3
                    )
                    if response.ok:
                        data = response.json()
                        if data.get("hidden_dim"):
                            arch = ModelArchitecture(
                                hidden_dim=data["hidden_dim"],
                                intermediate_dim=data.get("intermediate_dim", int(data["hidden_dim"] * 8 / 3)),
                                num_layers=data.get("num_layers", 12),
                                num_heads=data.get("num_heads", 12),
                                num_kv_heads=data.get("num_kv_heads", 4),
                            )
                            logger.debug(f"Got network architecture from peer {peer_url}: {arch.num_layers}L × {arch.hidden_dim}H")
                            return arch
                except Exception:
                    continue
        
        # Method 3: DHT lookup (if available)
        if self.p2p_manager and hasattr(self.p2p_manager, 'dht') and self.p2p_manager.dht:
            try:
                import hashlib
                key = int(hashlib.sha1("network_architecture".encode()).hexdigest(), 16)
                value = self.p2p_manager.dht.lookup_value(key)
                if value:
                    import json
                    data = json.loads(value)
                    if isinstance(data, dict) and data.get("hidden_dim"):
                        arch = ModelArchitecture(
                            hidden_dim=data["hidden_dim"],
                            intermediate_dim=data.get("intermediate_dim", int(data["hidden_dim"] * 8 / 3)),
                            num_layers=data.get("num_layers", 12),
                            num_heads=data.get("num_heads", 12),
                            num_kv_heads=data.get("num_kv_heads", 4),
                        )
                        logger.debug(f"Got network architecture from DHT: {arch.num_layers}L × {arch.hidden_dim}H")
                        return arch
            except Exception as e:
                logger.debug(f"DHT architecture lookup failed: {e}")
        
        return None
    
    def _architectures_compatible(self, arch1: ModelArchitecture, arch2: ModelArchitecture) -> bool:
        """
        Check if two architectures are compatible for gradient exchange.
        
        Compatible means: same hidden_dim, num_heads, num_kv_heads
        (num_layers can differ - nodes just hold different subsets)
        """
        return (
            arch1.hidden_dim == arch2.hidden_dim and
            arch1.num_heads == arch2.num_heads and
            arch1.num_kv_heads == arch2.num_kv_heads
        )
    
    def _archive_incompatible_checkpoint(self):
        """
        Archive an incompatible checkpoint instead of deleting it.
        
        Storage-aware: Keeps only MAX_ARCHIVED_CHECKPOINTS and respects
        the user's storage budget.
        """
        MAX_ARCHIVED_CHECKPOINTS = 2  # Keep at most 2 old checkpoints
        
        path = self.CHECKPOINT_DIR / f"dynamic_node_{self.wallet_id}.pt"
        
        if not path.exists():
            return
        
        # First, clean up old archives to stay within limits
        self._cleanup_old_archives(MAX_ARCHIVED_CHECKPOINTS - 1)  # Make room for new one
        
        # Now archive the current checkpoint
        import time
        timestamp = int(time.time())
        archive_path = self.CHECKPOINT_DIR / f"archived_{self.wallet_id}_{timestamp}.pt"
        
        try:
            path.rename(archive_path)
            logger.info(f"Archived incompatible checkpoint to: {archive_path.name}")
        except Exception as e:
            logger.warning(f"Could not archive checkpoint: {e}")
            # If archive fails, just delete it
            try:
                path.unlink()
                logger.info(f"Deleted incompatible checkpoint (archive failed)")
            except Exception:
                pass
    
    def _cleanup_old_archives(self, max_keep: int = 2):
        """
        Clean up old archived checkpoints, keeping only the most recent ones.
        
        Also enforces storage budget if archives are taking too much space.
        """
        # Find all archives for this wallet
        pattern = f"archived_{self.wallet_id}_*.pt"
        archives = sorted(
            self.CHECKPOINT_DIR.glob(pattern),
            key=lambda p: p.stat().st_mtime,
            reverse=True  # Newest first
        )
        
        # Calculate total archive size
        total_archive_mb = sum(p.stat().st_size / (1024 * 1024) for p in archives)
        
        # Storage budget: archives should use at most 20% of max_storage
        archive_budget_mb = self.max_storage_mb * 0.2
        
        # Delete archives that exceed count OR storage limits
        deleted_count = 0
        for i, archive in enumerate(archives):
            should_delete = False
            
            # Too many archives
            if i >= max_keep:
                should_delete = True
            
            # Over storage budget
            if total_archive_mb > archive_budget_mb:
                should_delete = True
            
            if should_delete:
                try:
                    archive_size_mb = archive.stat().st_size / (1024 * 1024)
                    archive.unlink()
                    total_archive_mb -= archive_size_mb
                    deleted_count += 1
                    logger.debug(f"Cleaned up old archive: {archive.name}")
                except Exception:
                    pass
        
        if deleted_count > 0:
            logger.info(f"Cleaned up {deleted_count} old archived checkpoint(s)")
    
    def _get_checkpoint_layer_count(self) -> Optional[int]:
        """
        Get the actual number of layers saved in checkpoint.
        
        This is important because the model may have GROWN beyond the base architecture.
        The architecture might say 11 layers, but 110 layers could be saved!
        """
        path = self.CHECKPOINT_DIR / f"dynamic_node_{self.wallet_id}.pt"
        legacy_path = self.CHECKPOINT_DIR / f"dynamic_node_{self.node_id[:16]}.pt"
        
        checkpoint_path = path if path.exists() else (legacy_path if legacy_path.exists() else None)
        if not checkpoint_path:
            return None
        
        try:
            checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
            layers = checkpoint.get("layers", {})
            if layers:
                return len(layers)
            # Fall back to layer_ids if present
            layer_ids = checkpoint.get("layer_ids", [])
            if layer_ids:
                return len(layer_ids)
        except Exception as e:
            logger.debug(f"Could not get checkpoint layer count: {e}")
        
        return None
    
    def _peek_checkpoint_architecture(self) -> Optional[ModelArchitecture]:
        """
        Peek at checkpoint to get saved architecture WITHOUT loading weights.

        This allows us to use the same architecture as the checkpoint,
        preventing architecture drift between restarts on the same machine.
        """
        # Use wallet_id for stable checkpoint path
        path = self.CHECKPOINT_DIR / f"dynamic_node_{self.wallet_id}.pt"
        
        # Also check legacy path
        legacy_path = self.CHECKPOINT_DIR / f"dynamic_node_{self.node_id[:16]}.pt"
        
        checkpoint_path = None
        if path.exists():
            checkpoint_path = path
        elif legacy_path.exists():
            checkpoint_path = legacy_path
        
        if not checkpoint_path:
            return None
        
        try:
            # Load just the metadata (weights_only would fail, but we catch it)
            checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
            
            arch_dict = checkpoint.get("architecture")
            if arch_dict:
                return ModelArchitecture.from_dict(arch_dict)
        except Exception as e:
            logger.debug(f"Could not peek checkpoint architecture: {e}")
        
        return None
    
    def _load_embedding_with_vocab_expansion(self, embedding: nn.Embedding, state_dict: dict, name: str):
        """
        Load embedding weights, handling vocab size expansion gracefully.
        
        When vocabulary grows (tokenizer learns new merges), the checkpoint has fewer
        tokens than the current model. This method:
        1. Loads weights for existing tokens (preserves all training)
        2. Keeps randomly initialized weights for new tokens
        """
        checkpoint_weight = state_dict.get("weight")
        if checkpoint_weight is None:
            logger.warning(f"[CHECKPOINT] No weight found in {name} state_dict")
            return
        
        checkpoint_vocab_size = checkpoint_weight.shape[0]
        current_vocab_size = embedding.weight.shape[0]
        
        if checkpoint_vocab_size == current_vocab_size:
            # Same size - normal load
            embedding.load_state_dict(state_dict)
            logger.info(f"[CHECKPOINT] Loaded {name}: {current_vocab_size} tokens")
        elif checkpoint_vocab_size < current_vocab_size:
            # Vocab expanded - partial load (PRESERVE TRAINING!)
            with torch.no_grad():
                embedding.weight[:checkpoint_vocab_size] = checkpoint_weight
            logger.info(f"[CHECKPOINT] Loaded {name} with vocab expansion: "
                       f"{checkpoint_vocab_size} → {current_vocab_size} tokens "
                       f"(preserved {checkpoint_vocab_size} trained embeddings)")
        else:
            # Vocab shrunk (unusual) - load what fits
            with torch.no_grad():
                embedding.weight[:] = checkpoint_weight[:current_vocab_size]
            logger.warning(f"[CHECKPOINT] Loaded {name} with vocab truncation: "
                          f"{checkpoint_vocab_size} → {current_vocab_size} tokens")
    
    def _load_lm_head_with_vocab_expansion(self, lm_head: nn.Linear, state_dict: dict, name: str):
        """
        Load LM head weights, handling vocab size expansion gracefully.
        
        Similar to embedding expansion - preserves trained weights for existing tokens.
        """
        checkpoint_weight = state_dict.get("weight")
        checkpoint_bias = state_dict.get("bias")
        
        if checkpoint_weight is None:
            logger.warning(f"[CHECKPOINT] No weight found in {name} state_dict")
            return
        
        checkpoint_vocab_size = checkpoint_weight.shape[0]
        current_vocab_size = lm_head.weight.shape[0]
        
        if checkpoint_vocab_size == current_vocab_size:
            # Same size - normal load
            lm_head.load_state_dict(state_dict)
            logger.info(f"[CHECKPOINT] Loaded {name}: {current_vocab_size} outputs")
        elif checkpoint_vocab_size < current_vocab_size:
            # Vocab expanded - partial load (PRESERVE TRAINING!)
            with torch.no_grad():
                lm_head.weight[:checkpoint_vocab_size] = checkpoint_weight
                if checkpoint_bias is not None and lm_head.bias is not None:
                    lm_head.bias[:checkpoint_vocab_size] = checkpoint_bias
            logger.info(f"[CHECKPOINT] Loaded {name} with vocab expansion: "
                       f"{checkpoint_vocab_size} → {current_vocab_size} outputs "
                       f"(preserved {checkpoint_vocab_size} trained weights)")
        else:
            # Vocab shrunk (unusual) - load what fits
            with torch.no_grad():
                lm_head.weight[:] = checkpoint_weight[:current_vocab_size]
                if checkpoint_bias is not None and lm_head.bias is not None:
                    lm_head.bias[:] = checkpoint_bias[:current_vocab_size]
            logger.warning(f"[CHECKPOINT] Loaded {name} with vocab truncation: "
                          f"{checkpoint_vocab_size} → {current_vocab_size} outputs")
    
    def _load_checkpoint(self):
        """Load checkpoint from disk if it exists (resume training)."""
        # Use wallet_id for stable checkpoint path (survives node_id changes)
        path = self.CHECKPOINT_DIR / f"dynamic_node_{self.wallet_id}.pt"
        
        # Also check legacy path (node_id-based) for migration
        legacy_path = self.CHECKPOINT_DIR / f"dynamic_node_{self.node_id[:16]}.pt"
        if not path.exists() and legacy_path.exists():
            logger.info(f"Migrating checkpoint from legacy path: {legacy_path.name} -> {path.name}")
            legacy_path.rename(path)
        
        if not path.exists():
            logger.info(f"No checkpoint found at {path.name}, starting fresh")
            return False
        
        logger.info(f"Loading checkpoint from: {path.name}")
        try:
            checkpoint = torch.load(path, map_location=self.device, weights_only=False)
            
            # ARCHITECTURE COMPATIBILITY CHECK
            saved_arch_dict = checkpoint.get("architecture")
            if saved_arch_dict:
                saved_arch = ModelArchitecture.from_dict(saved_arch_dict)
                current_arch = self.model.architecture
                
                # Check if architecture changed (includes num_heads for head_dim compatibility)
                if (saved_arch.hidden_dim != current_arch.hidden_dim or 
                    saved_arch.intermediate_dim != current_arch.intermediate_dim or
                    saved_arch.num_heads != current_arch.num_heads or
                    saved_arch.num_kv_heads != current_arch.num_kv_heads):
                    logger.warning(f"Architecture mismatch! Checkpoint is incompatible.")
                    logger.warning(f"  Saved: {saved_arch.num_layers}L × {saved_arch.hidden_dim}H, "
                                   f"heads={saved_arch.num_heads}/{saved_arch.num_kv_heads}")
                    logger.warning(f"  Current: {current_arch.num_layers}L × {current_arch.hidden_dim}H, "
                                   f"heads={current_arch.num_heads}/{current_arch.num_kv_heads}")
                    logger.warning(f"  Starting fresh (architecture was upgraded)")
                    # Delete incompatible checkpoint
                    try:
                        path.unlink()
                        logger.info(f"Deleted incompatible checkpoint: {path}")
                    except Exception:
                        pass
                    return False
            else:
                logger.warning("Legacy checkpoint without architecture info - starting fresh")
                # Delete legacy checkpoint
                try:
                    path.unlink()
                    logger.info(f"Deleted legacy checkpoint: {path}")
                except Exception:
                    pass
                return False
            
            # Check layer assignment compatibility
            saved_layers = set(checkpoint.get("layer_ids", []))
            current_layers = set(self.my_layer_ids)
            
            if saved_layers != current_layers:
                # Layers changed - try to load what we can
                common_layers = saved_layers.intersection(current_layers)
                if common_layers:
                    logger.warning(f"Layer assignment changed: saved={len(saved_layers)}, current={len(current_layers)}, common={len(common_layers)}")
                    logger.info(f"Will load {len(common_layers)} common layers from checkpoint")
                else:
                    logger.warning(f"No common layers between checkpoint and current assignment, starting fresh")
                    return False
            
            # Load layer weights
            for layer_id, state_dict in checkpoint.get("layers", {}).items():
                layer_id = int(layer_id)
                if layer_id in self.model.my_layers:
                    self.model.my_layers[layer_id].load_state_dict(state_dict)
            
            # Load embedding if present (handle vocab size changes gracefully)
            if self.model.embedding and "embedding" in checkpoint:
                self._load_embedding_with_vocab_expansion(
                    self.model.embedding, 
                    checkpoint["embedding"],
                    "embedding"
                )
            
            # Load LM head if present (handle vocab size changes gracefully)
            if self.model.lm_head and "lm_head" in checkpoint:
                self._load_lm_head_with_vocab_expansion(
                    self.model.lm_head,
                    checkpoint["lm_head"],
                    "lm_head"
                )
            
            # Load final norm if present
            if self.model.final_norm and "final_norm" in checkpoint:
                self.model.final_norm.load_state_dict(checkpoint["final_norm"])
            
            # Restore training state
            self.total_training_rounds = checkpoint.get("total_training_rounds", 0)
            
            # Store optimizer state for later loading (after optimizer is created)
            if "optimizer" in checkpoint:
                self._pending_optimizer_state = checkpoint["optimizer"]
            
            # Store DiLoCo state for later loading (after swarm is created)
            if "diloco" in checkpoint:
                self._pending_diloco_state = checkpoint["diloco"]
                logger.info("[NODE] DiLoCo state found in checkpoint, will restore after swarm init")
            
            # Count how many layers were actually loaded
            loaded_layer_count = sum(1 for lid in checkpoint.get("layers", {}).keys() if int(lid) in self.model.my_layers)
            logger.info(f"Checkpoint loaded: {self.total_training_rounds} training rounds, "
                       f"{loaded_layer_count}/{len(current_layers)} layers from {path}")
            return True
            
        except Exception as e:
            logger.warning(f"Failed to load checkpoint: {e}, starting fresh")
            return False
    
    def _restore_pending_state(self):
        """
        Restore optimizer and DiLoCo state after they are initialized.
        
        Called after swarm/optimizer are set up to restore checkpoint state.
        """
        # Restore optimizer state
        if hasattr(self, '_pending_optimizer_state') and self._pending_optimizer_state:
            if hasattr(self, 'optimizer') and self.optimizer:
                try:
                    self.optimizer.load_state_dict(self._pending_optimizer_state)
                    logger.info("[NODE] Restored optimizer state from checkpoint")
                except Exception as e:
                    logger.warning(f"[NODE] Could not restore optimizer state: {e}")
            del self._pending_optimizer_state
        
        # Restore DiLoCo state
        if hasattr(self, '_pending_diloco_state') and self._pending_diloco_state:
            if hasattr(self, 'swarm') and self.swarm:
                diloco = getattr(self.swarm, 'diloco_trainer', None)
                if diloco and hasattr(diloco, 'load_state_dict'):
                    try:
                        diloco.load_state_dict(self._pending_diloco_state)
                        logger.info(f"[NODE] Restored DiLoCo state (inner_step={diloco.stats.inner_step_count})")
                    except Exception as e:
                        logger.warning(f"[NODE] Could not restore DiLoCo state: {e}")
            del self._pending_diloco_state
    
    # Class-level save lock to prevent concurrent checkpoint saves
    _checkpoint_save_lock = threading.Lock()
    _checkpoint_save_in_progress = False
    
    def _save_checkpoint(self, async_save: bool = True):
        """
        Smart checkpoint saving with STREAMING ASYNC for memory-constrained systems.
        
        THREE MODES:
        1. BULK ASYNC (>32GB free OR >2.5x checkpoint): Clone all, save in thread
        2. STREAMING ASYNC (>500MB free): Clone one layer at a time, save incrementally
        3. SYNC (<500MB free): Blocking save (last resort)
        
        Streaming async blocks training only during the brief snapshot (~1-2s for 110 layers),
        then saves to disk in background (~10-60s) while training continues.
        
        Thread-safe: concurrent saves are serialized via lock.
        """
        if not self.model:
            return
        
        # Prevent concurrent saves (async save might still be in progress)
        if not DynamicNeuroNode._checkpoint_save_lock.acquire(blocking=False):
            logger.debug("[NODE] Checkpoint save skipped - another save in progress")
            return
        
        # Use wallet_id for stable checkpoint path
        path = self.CHECKPOINT_DIR / f"dynamic_node_{self.wallet_id}.pt"
        temp_path = self.CHECKPOINT_DIR / f"dynamic_node_{self.wallet_id}.pt.tmp"

        # 1. Assess memory situation (respect configured --memory limit!)
        try:
            total_params = sum(p.numel() for p in self.model.parameters())
            checkpoint_size_mb = (total_params * 4) / (1024 * 1024)
            
            # Get ACTUAL available memory (respecting configured limit)
            vm = psutil.virtual_memory()
            system_available_mb = vm.available / (1024 * 1024)
            
            # Check current process memory usage
            process = psutil.Process()
            process_used_mb = process.memory_info().rss / (1024 * 1024)
            
            # If user set --memory limit, respect it
            # Available = min(system_available, configured_limit - current_usage)
            configured_limit = getattr(self, 'available_memory_mb', None)
            if configured_limit:
                # How much headroom do we have within our configured limit?
                headroom_mb = max(0, configured_limit - process_used_mb)
                # Use the more conservative of system available or our headroom
                available_mb = min(system_available_mb, headroom_mb)
                logger.debug(f"[NODE] Memory check: process={process_used_mb:.0f}MB, "
                           f"limit={configured_limit:.0f}MB, headroom={headroom_mb:.0f}MB, "
                           f"system_free={system_available_mb:.0f}MB, using={available_mb:.0f}MB")
            else:
                available_mb = system_available_mb
            
            # Determine save mode based on available headroom
            # Bulk async needs 2.5x checkpoint size to clone everything
            can_bulk_async = (available_mb > (checkpoint_size_mb * 2.5)) or (available_mb > 32000)
            # Streaming async just needs enough for 1 layer (~50-100MB typically)
            can_stream_async = available_mb > 500
            
        except Exception as e:
            logger.warning(f"[NODE] Could not assess memory: {e}. Using streaming async.")
            can_bulk_async = False
            can_stream_async = True
            checkpoint_size_mb = 0
            available_mb = 0
        
        try:
            # ============ ATOMIC SNAPSHOT PHASE ============
            # ALL state must be captured together to ensure consistency.
            # DiLoCo state + model weights must be from the same "moment in time".
            
            # Helper to deep-clone DiLoCo state (ALL tensors to CPU for async safety)
            def _clone_diloco_state():
                if not hasattr(self, 'swarm') or not self.swarm:
                    return None
                diloco = getattr(self.swarm, 'diloco_trainer', None)
                if not diloco or not hasattr(diloco, 'state_dict'):
                    return None
                try:
                    state = diloco.state_dict()
                    
                    # Deep clone optimizer state (handles both PyTorch and custom formats)
                    def _clone_optimizer_state(opt_state):
                        if opt_state is None:
                            return None
                        cloned = {}
                        for key, value in opt_state.items():
                            if isinstance(value, torch.Tensor):
                                # Direct tensor (e.g., in custom optimizers)
                                cloned[key] = value.detach().clone().cpu()
                            elif isinstance(value, dict):
                                # Nested dict (e.g., 'state' or 'velocity' dicts)
                                cloned[key] = {}
                                for k, v in value.items():
                                    if isinstance(v, torch.Tensor):
                                        cloned[key][k] = v.detach().clone().cpu()
                                    elif isinstance(v, dict):
                                        # PyTorch optimizer 'state' has param_idx -> {key: tensor}
                                        cloned[key][k] = {}
                                        for kk, vv in v.items():
                                            if isinstance(vv, torch.Tensor):
                                                cloned[key][k][kk] = vv.detach().clone().cpu()
                                            else:
                                                cloned[key][k][kk] = vv
                                    else:
                                        cloned[key][k] = v
                            elif isinstance(value, list):
                                # List (e.g., param_groups) - shallow copy is fine
                                cloned[key] = list(value)
                            else:
                                # Scalar values (lr, momentum, etc.)
                                cloned[key] = value
                        return cloned
                    
                    # Deep clone all tensors to CPU for async safety
                    cloned = {
                        'config': dict(state.get('config', {})),
                        'inner_optimizer': _clone_optimizer_state(state.get('inner_optimizer')),
                        'outer_optimizer': _clone_optimizer_state(state.get('outer_optimizer')),
                        'initial_weights': {
                            k: v.detach().clone().cpu() 
                            for k, v in state.get('initial_weights', {}).items()
                        },
                        'stats': dict(state.get('stats', {})),
                        'phase': state.get('phase', 'idle'),
                    }
                    return cloned
                except Exception as e:
                    logger.warning(f"[NODE] Could not snapshot DiLoCo state: {e}")
                    return None
            
            # ============ MODE 1: BULK ASYNC (plenty of memory) ============
            if async_save and can_bulk_async:
                logger.debug(f"[NODE] Checkpoint: BULK ASYNC (Free: {available_mb:.0f}MB)")
                
                # Capture everything atomically
                checkpoint = {
                    "node_id": self.node_id,
                    "layer_ids": list(self.my_layer_ids),
                    "architecture": self.model.architecture.to_dict(),
                    "architecture_version": self.layer_pool.architecture_version if self.layer_pool else 1,
                    "has_embedding": self.model.has_embedding,
                    "has_lm_head": self.model.has_lm_head,
                    "total_training_rounds": self.total_training_rounds,
                    "current_loss": self.current_loss,
                    "timestamp": time.time(),
                    "layers": {
                        layer_id: {k: v.clone().cpu() for k, v in layer.state_dict().items()}
                        for layer_id, layer in self.model.my_layers.items()
                    },
                }
                if self.model.embedding:
                    checkpoint["embedding"] = {k: v.clone().cpu() for k, v in self.model.embedding.state_dict().items()}
                if self.model.lm_head:
                    checkpoint["lm_head"] = {k: v.clone().cpu() for k, v in self.model.lm_head.state_dict().items()}
                if self.model.final_norm:
                    checkpoint["final_norm"] = {k: v.clone().cpu() for k, v in self.model.final_norm.state_dict().items()}
                
                # DiLoCo state captured AFTER model weights (both in same atomic snapshot)
                diloco_state = _clone_diloco_state()
                if diloco_state:
                    checkpoint["diloco"] = diloco_state
                
                def _do_bulk_save():
                    try:
                        torch.save(checkpoint, temp_path, _use_new_zipfile_serialization=False)
                        import shutil
                        shutil.move(str(temp_path), str(path))
                        logger.info(f"[NODE] Checkpoint saved ({len(self.my_layer_ids)} layers)")
                    except Exception as e:
                        logger.error(f"[NODE] Checkpoint save failed: {e}")
                        if temp_path.exists(): temp_path.unlink()
                    finally:
                        DynamicNeuroNode._checkpoint_save_lock.release()
                
                # Use daemon=False so checkpoint completes even during shutdown
                threading.Thread(target=_do_bulk_save, daemon=False).start()
                return  # Lock will be released by background thread
            
            # ============ MODE 2: STREAMING ASYNC (memory-efficient) ============
            if async_save and can_stream_async:
                logger.info(f"[NODE] Checkpoint: STREAMING ASYNC (Free: {available_mb:.0f}MB, cloning {len(self.model.my_layers)} layers)")
                
                # SNAPSHOT PHASE: Clone one layer at a time into a list
                # This brief pause (~1-2s) ensures consistency without needing full clone memory
                snapshot_start = time.time()
                
                # Capture metadata first (lightweight)
                checkpoint_meta = {
                    "node_id": self.node_id,
                    "layer_ids": list(self.my_layer_ids),
                    "architecture": self.model.architecture.to_dict(),
                    "architecture_version": self.layer_pool.architecture_version if self.layer_pool else 1,
                    "has_embedding": self.model.has_embedding,
                    "has_lm_head": self.model.has_lm_head,
                    "total_training_rounds": self.total_training_rounds,
                    "current_loss": self.current_loss,
                    "timestamp": time.time(),
                }
                
                # Clone layers one at a time (memory efficient)
                layer_snapshots = []
                for layer_id, layer in self.model.my_layers.items():
                    layer_state = {k: v.detach().clone().cpu() for k, v in layer.state_dict().items()}
                    layer_snapshots.append((layer_id, layer_state))
                
                # Clone special modules
                embedding_snapshot = None
                lm_head_snapshot = None
                final_norm_snapshot = None
                
                if self.model.embedding:
                    embedding_snapshot = {k: v.detach().clone().cpu() for k, v in self.model.embedding.state_dict().items()}
                if self.model.lm_head:
                    lm_head_snapshot = {k: v.detach().clone().cpu() for k, v in self.model.lm_head.state_dict().items()}
                if self.model.final_norm:
                    final_norm_snapshot = {k: v.detach().clone().cpu() for k, v in self.model.final_norm.state_dict().items()}
                
                # DiLoCo state - captured in SAME snapshot window as model weights
                diloco_snapshot = _clone_diloco_state()
                
                snapshot_time = time.time() - snapshot_start
                logger.debug(f"[NODE] Snapshot complete in {snapshot_time:.1f}s, starting async save")
                
                # ASYNC SAVE PHASE: Write to disk in background thread
                # All data is now cloned and owned by this closure - safe for async
                def _do_streaming_save():
                    try:
                        checkpoint = dict(checkpoint_meta)
                        checkpoint["layers"] = {lid: lstate for lid, lstate in layer_snapshots}
                        if embedding_snapshot:
                            checkpoint["embedding"] = embedding_snapshot
                        if lm_head_snapshot:
                            checkpoint["lm_head"] = lm_head_snapshot
                        if final_norm_snapshot:
                            checkpoint["final_norm"] = final_norm_snapshot
                        if diloco_snapshot:
                            checkpoint["diloco"] = diloco_snapshot
                        
                        torch.save(checkpoint, temp_path, _use_new_zipfile_serialization=False)
                        import shutil
                        shutil.move(str(temp_path), str(path))
                        logger.info(f"[NODE] Checkpoint saved ({len(layer_snapshots)} layers)")
                    except Exception as e:
                        logger.error(f"[NODE] Checkpoint save failed: {e}")
                        if temp_path.exists(): temp_path.unlink()
                    finally:
                        DynamicNeuroNode._checkpoint_save_lock.release()
                
                # Use daemon=False so checkpoint completes even during shutdown
                threading.Thread(target=_do_streaming_save, daemon=False).start()
                return  # Lock will be released by background thread
            
            # ============ MODE 3: SYNC (last resort, very low memory) ============
            logger.warning(f"[NODE] Checkpoint: SYNC mode (Free: {available_mb:.0f}MB < 500MB minimum)")
            
            checkpoint = {
                "node_id": self.node_id,
                "layer_ids": list(self.my_layer_ids),
                "architecture": self.model.architecture.to_dict(),
                "architecture_version": self.layer_pool.architecture_version if self.layer_pool else 1,
                "has_embedding": self.model.has_embedding,
                "has_lm_head": self.model.has_lm_head,
                "total_training_rounds": self.total_training_rounds,
                "current_loss": self.current_loss,
                "timestamp": time.time(),
                "layers": {
                    layer_id: layer.state_dict()
                    for layer_id, layer in self.model.my_layers.items()
                },
            }
            if self.model.embedding:
                checkpoint["embedding"] = self.model.embedding.state_dict()
            if self.model.lm_head:
                checkpoint["lm_head"] = self.model.lm_head.state_dict()
            if self.model.final_norm:
                checkpoint["final_norm"] = self.model.final_norm.state_dict()
            
            # DiLoCo state (no need to clone for sync - we block anyway)
            diloco_state = _clone_diloco_state()
            if diloco_state:
                checkpoint["diloco"] = diloco_state
            
            torch.save(checkpoint, temp_path, _use_new_zipfile_serialization=False)
            import shutil
            shutil.move(str(temp_path), str(path))
            logger.info(f"[NODE] Checkpoint saved ({len(self.my_layer_ids)} layers)")
            
            # Sync mode completed successfully, release lock
            DynamicNeuroNode._checkpoint_save_lock.release()
            
        except Exception as e:
            logger.error(f"[NODE] Checkpoint preparation failed: {type(e).__name__}: {e}")
            try:
                if temp_path.exists(): temp_path.unlink()
            except:
                pass
            # Release lock on exception
            DynamicNeuroNode._checkpoint_save_lock.release()


def create_dynamic_node(
    node_token: str,
    port: int = 8000,
    tracker_url: str = "https://neuroshard.com/api/tracker",
    available_memory_mb: Optional[float] = None,
    enable_training: bool = True,
    max_storage_mb: float = 100.0,
    max_cpu_threads: Optional[int] = None,
    device: str = "auto",
    p2p_manager: Optional[Any] = None,  # NEW: Pass P2P for DHT discovery during layer assignment
) -> DynamicNeuroNode:
    """
    Create and start a dynamic node.

    MULTI-NODE SUPPORT:
    If the same token is used on multiple machines or ports, each gets a unique
    node_id (based on machine + port) while sharing the same wallet_id (based on token).

    This means:
    - Each physical node has a unique network identity
    - Earnings accumulate to the same NEURO wallet
    - No conflicts in DHT/layer assignments
    
    FULLY DECENTRALIZED:
    If p2p_manager is provided, DHT is used for network discovery during layer
    assignment. No tracker fallbacks - pure P2P!
    """
    from neuroshard.utils.hardware import get_instance_id

    # Generate instance-specific node_id
    instance_id = get_instance_id(port)

    # Combine token with instance for unique network identity
    # wallet_id (from token alone) is used for NEURO earnings
    # node_id (from token + instance) is used for network identity
    combined = f"{node_token}:{instance_id}"
    node_id = str(int(hashlib.sha256(combined.encode()).hexdigest(), 16))
    
    # Log multi-node info
    wallet_id = hashlib.sha256(node_token.encode()).hexdigest()[:16]
    logger.info(f"Instance ID: {instance_id} (machine+port)")
    logger.info(f"Wallet ID: {wallet_id}... (for NEURO earnings)")
    logger.info(f"Node ID: {node_id[:16]}... (unique network identity)")
    
    node = DynamicNeuroNode(
        node_id=node_id,
        port=port,
        tracker_url=tracker_url,
        node_token=node_token,
        available_memory_mb=available_memory_mb,
        enable_training=enable_training,
        max_storage_mb=max_storage_mb,
        max_cpu_threads=max_cpu_threads,
        device=device,
    )
    
    # Store instance info for debugging
    node.instance_id = instance_id
    node.wallet_id = wallet_id
    
    # CRITICAL: Connect P2P BEFORE start() so DHT is available for layer discovery!
    # This enables fully decentralized network discovery without tracker fallbacks.
    if p2p_manager:
        node.p2p_manager = p2p_manager
        logger.info("P2P connected BEFORE start - DHT available for network discovery")
    
    node.start()
    
    return node

