File size: 5,032 Bytes
5c93746 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """
Data containers for communication operations.
This module defines data structures used for communication between distributed ranks.
"""
from dataclasses import dataclass
from typing import Optional, List, Dict, Any
import torch
@dataclass
class LatentData:
"""
Container for latent data and related information.
This class encapsulates all the data that needs to be transferred between ranks
during the inference pipeline.
"""
chunk_idx: int
latents: torch.Tensor
original_latents: torch.Tensor
current_start: torch.Tensor
current_end: torch.Tensor
current_step: int
patched_x_shape: torch.Tensor
def __post_init__(self):
"""Validate tensor shapes and types after initialization."""
if not isinstance(self.latents, torch.Tensor):
raise TypeError("latents must be a torch.Tensor")
if not isinstance(self.original_latents, torch.Tensor):
raise TypeError("original_latents must be a torch.Tensor")
if not isinstance(self.current_start, torch.Tensor):
raise TypeError("current_start must be a torch.Tensor")
if not isinstance(self.current_end, torch.Tensor):
raise TypeError("current_end must be a torch.Tensor")
if not isinstance(self.patched_x_shape, torch.Tensor):
raise TypeError("patched_x_shape must be a torch.Tensor")
@dataclass
class KVCacheData:
"""
Container for KV cache data.
This class encapsulates key-value cache information for transformer blocks.
"""
block_index: int
k_cache: torch.Tensor
v_cache: torch.Tensor
global_end_index: torch.Tensor
local_end_index: torch.Tensor
def __post_init__(self):
"""Validate tensor shapes and types after initialization."""
if not isinstance(self.k_cache, torch.Tensor):
raise TypeError("k_cache must be a torch.Tensor")
if not isinstance(self.v_cache, torch.Tensor):
raise TypeError("v_cache must be a torch.Tensor")
if not isinstance(self.global_end_index, torch.Tensor):
raise TypeError("global_end_index must be a torch.Tensor")
if not isinstance(self.local_end_index, torch.Tensor):
raise TypeError("local_end_index must be a torch.Tensor")
@dataclass
class CommunicationConfig:
"""
Configuration for communication operations.
This class holds configuration parameters for distributed communication.
"""
max_outstanding: int = 1
buffer_pool_size: int = 10
enable_buffer_reuse: bool = True
communication_timeout: float = 30.0
enable_async_communication: bool = True
def __post_init__(self):
"""Validate configuration parameters."""
if self.max_outstanding < 1:
raise ValueError("max_outstanding must be at least 1")
if self.buffer_pool_size < 1:
raise ValueError("buffer_pool_size must be at least 1")
if self.communication_timeout <= 0:
raise ValueError("communication_timeout must be positive")
@dataclass
class BlockInterval:
"""
Container for block interval information.
This class represents a block interval [start, end) for a specific rank.
"""
start: int
end: int
rank: int
def __post_init__(self):
"""Validate block interval parameters."""
if self.start < 0:
raise ValueError("start must be non-negative")
if self.end <= self.start:
raise ValueError("end must be greater than start")
if self.rank < 0:
raise ValueError("rank must be non-negative")
@property
def size(self) -> int:
"""Get the size of the block interval."""
return self.end - self.start
def contains(self, block_index: int) -> bool:
"""Check if the block interval contains the given block index."""
return self.start <= block_index < self.end
@dataclass
class PerformanceMetrics:
"""
Container for performance metrics.
This class holds timing and performance information for communication operations.
"""
dit_time: float
total_time: float
communication_time: float
buffer_allocation_time: float
def __post_init__(self):
"""Validate performance metrics."""
if self.dit_time < 0:
raise ValueError("dit_time must be non-negative")
if self.total_time < 0:
raise ValueError("total_time must be non-negative")
if self.communication_time < 0:
raise ValueError("communication_time must be non-negative")
if self.buffer_allocation_time < 0:
raise ValueError("buffer_allocation_time must be non-negative")
@property
def efficiency(self) -> float:
"""Calculate communication efficiency (computation time / total time)."""
if self.total_time == 0:
return 0.0
return (self.total_time - self.communication_time) / self.total_time
|