|
|
"""
|
|
|
Tensor operations for distributed computing.
|
|
|
"""
|
|
|
import torch
|
|
|
import numpy as np
|
|
|
from typing import Dict, List, Optional, Union, Tuple
|
|
|
|
|
|
class TensorOps:
|
|
|
"""Utility class for distributed tensor operations."""
|
|
|
|
|
|
@staticmethod
|
|
|
def split_tensor(tensor: torch.Tensor, num_parts: int) -> List[torch.Tensor]:
|
|
|
"""Split a tensor into multiple parts for distributed processing."""
|
|
|
return torch.chunk(tensor, num_parts)
|
|
|
|
|
|
@staticmethod
|
|
|
def merge_tensors(tensors: List[torch.Tensor], dim: int = 0) -> torch.Tensor:
|
|
|
"""Merge multiple tensors back into a single tensor."""
|
|
|
return torch.cat(tensors, dim=dim)
|
|
|
|
|
|
@staticmethod
|
|
|
def average_gradients(gradients: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
|
|
|
"""Average gradients from multiple workers."""
|
|
|
avg_gradients = {}
|
|
|
for key in gradients[0].keys():
|
|
|
avg_gradients[key] = torch.mean(torch.stack([g[key] for g in gradients]), dim=0)
|
|
|
return avg_gradients
|
|
|
|
|
|
@staticmethod
|
|
|
def serialize_tensor(tensor: torch.Tensor) -> Dict[str, Union[List, str]]:
|
|
|
"""Serialize a tensor for storage/transmission."""
|
|
|
return {
|
|
|
'data': tensor.cpu().numpy().tolist(),
|
|
|
'shape': list(tensor.shape),
|
|
|
'dtype': str(tensor.dtype)
|
|
|
}
|
|
|
|
|
|
@staticmethod
|
|
|
def deserialize_tensor(tensor_dict: Dict[str, Union[List, str]]) -> torch.Tensor:
|
|
|
"""Deserialize a tensor from storage/transmission format."""
|
|
|
data = np.array(tensor_dict['data'])
|
|
|
shape = tensor_dict['shape']
|
|
|
dtype = getattr(torch, tensor_dict['dtype'].split('.')[-1])
|
|
|
return torch.tensor(data, dtype=dtype).reshape(shape)
|
|
|
|
|
|
@staticmethod
|
|
|
def gradient_clipping(gradients: Dict[str, torch.Tensor], max_norm: float) -> Dict[str, torch.Tensor]:
|
|
|
"""Apply gradient clipping to prevent exploding gradients."""
|
|
|
for k, v in gradients.items():
|
|
|
if v is not None:
|
|
|
torch.nn.utils.clip_grad_norm_(v, max_norm)
|
|
|
return gradients
|
|
|
|
|
|
@staticmethod
|
|
|
def reduce_precision(tensor: torch.Tensor, bits: int = 16) -> torch.Tensor:
|
|
|
"""Reduce tensor precision for efficient transmission."""
|
|
|
if bits == 16:
|
|
|
return tensor.half()
|
|
|
elif bits == 32:
|
|
|
return tensor.float()
|
|
|
else:
|
|
|
raise ValueError("Unsupported precision bits")
|
|
|
|
|
|
@staticmethod
|
|
|
def shard_tensor(tensor: torch.Tensor, shard_size: int) -> List[torch.Tensor]:
|
|
|
"""Shard a tensor into smaller pieces for distributed processing."""
|
|
|
return [tensor[i:i + shard_size] for i in range(0, tensor.size(0), shard_size)]
|
|
|
|
|
|
@staticmethod
|
|
|
def compute_parameter_norm(parameters: Dict[str, torch.Tensor]) -> float:
|
|
|
"""Compute the total norm of all parameters."""
|
|
|
total_norm = 0.0
|
|
|
for param in parameters.values():
|
|
|
total_norm += param.norm().item() ** 2
|
|
|
return total_norm ** 0.5 |
