DeepSeek-V3/inference/models/linear.py
Hitesh Yadav bc9459df40 refactor(inference): modularize model architecture for improved maintainability
BREAKING CHANGE: Restructured model.py into dedicated modules under inference/models/

Key Changes:
- Split monolithic model.py into focused, single-responsibility modules:
  - config.py: Model configuration and hyperparameters
  - attention.py: Multi-head Latent Attention (MLA) implementation
  - moe.py: Mixture of Experts components (Gate, Expert, MoE)
  - linear.py: Linear layer variants with parallel processing support
  - __init__.py: Clean public API exports

Benefits:
- Improved code organization and maintainability
- Better separation of concerns
- Enhanced testability of individual components
- Clearer dependency management
- Simplified future modifications and extensions

Migration:
- Update imports to use new module structure
- No functional changes to existing implementations
- Backwards compatible with current model weights
2025-01-05 16:28:10 +05:30

28 lines
1.0 KiB
Python

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from ..kernel import act_quant, weight_dequant, fp8_gemm
class Linear(nn.Module):
dtype = torch.bfloat16
def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
# ... (Linear implementation)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# ... (Linear forward implementation)
class ColumnParallelLinear(Linear):
def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
# ... (ColumnParallelLinear implementation)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# ... (ColumnParallelLinear forward implementation)
class RowParallelLinear(Linear):
def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):
# ... (RowParallelLinear implementation)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# ... (RowParallelLinear forward implementation)