Update core/data_architecture/sparse_ops.py

core/data_architecture/sparse_ops.py

@@ -18,10 +18,13 @@ import apex
 from apex import amp
 from apex.optimizers import FusedAdam
 
+# Assuming fused_ops is compiled and available
+import fused_ops  # Custom CUDA extension from fused_ops.cu
+
 class SparseLinear(nn.Module):
     """
     Sparse Linear Layer with Tensor Core Optimizations and Dynamic Pruning.
-    Prunes weights based on magnitude to improve efficiency on GPU.
+    Integrates fused GEMM + ReLU CUDA kernel for GPU efficiency.
     """
     def __init__(self, in_features, out_features, sparsity=0.5, use_fp16=True, dynamic_pruning=False):
         super(SparseLinear, self).__init__()
@@ -29,7 +32,7 @@ class SparseLinear(nn.Module):
         self.out_features = out_features
         self.sparsity = sparsity
         self.use_fp16 = use_fp16
-        self.dynamic_pruning = dynamic_pruning  # Toggle dynamic vs static pruning
+        self.dynamic_pruning = dynamic_pruning
 
         # Initialize dense weight and bias
         self.weight = nn.Parameter(
@@ -39,7 +42,7 @@ class SparseLinear(nn.Module):
             torch.zeros(out_features, dtype=torch.float16 if use_fp16 else torch.float32)
         )
 
-        # Sparse mask (static unless dynamic_pruning is enabled)
+        # Sparse mask
         self.register_buffer("mask", self.generate_mask())
 
     def generate_mask(self):
@@ -47,30 +50,29 @@ class SparseLinear(nn.Module):
         Generates a binary mask based on weight magnitude for structured sparsity.
         """
         if self.dynamic_pruning:
-            # Dynamic pruning will recompute this in forward pass
             return torch.ones_like(self.weight)
         weights_abs = self.weight.abs()
         threshold = torch.quantile(weights_abs.flatten(), self.sparsity)
-        return (weights_abs > threshold).float()
+        return (weights_abs > threshold).to(self.weight.dtype)
 
     def update_mask(self):
         """Update mask dynamically based on current weight magnitudes."""
         if self.dynamic_pruning:
             weights_abs = self.weight.abs()
             threshold = torch.quantile(weights_abs.flatten(), self.sparsity)
-            self.mask.data = (weights_abs > threshold).float()
+            self.mask.data = (weights_abs > threshold).to(self.weight.dtype)
 
     def forward(self, x):
         if self.dynamic_pruning:
-            self.update_mask()  # Recompute mask if dynamic pruning is enabled
+            self.update_mask()
 
-        if self.use_fp16:
-            with autocast():
-                pruned_weight = self.weight * self.mask
-                return F.linear(x, pruned_weight, self.bias)
+        if self.use_fp16 and x.is_cuda():
+            # Use fused CUDA kernel for GEMM + ReLU
+            return fused_ops.fused_sparse_gemm_relu(x, self.weight, self.mask, self.bias)
         else:
+            # Fallback to PyTorch
             pruned_weight = self.weight.float() * self.mask.float()
-            return F.linear(x.float(), pruned_weight, self.bias.float())
+            return F.relu(F.linear(x.float(), pruned_weight, self.bias.float()))
 
 
 class SparseConv2d(nn.Module):
@@ -84,7 +86,7 @@ class SparseConv2d(nn.Module):
         self.use_fp16 = use_fp16
         self.sparsity = sparsity
         self.dynamic_pruning = dynamic_pruning
-        self.block_size = block_size  # Optional block sparsity (e.g., (2, 2))
+        self.block_size = block_size
 
         self.conv = nn.Conv2d(
             in_channels,
@@ -97,16 +99,11 @@ class SparseConv2d(nn.Module):
         self.register_buffer("mask", self.generate_mask())
 
     def generate_mask(self):
-        """
-        Generate a mask based on weight magnitude, optionally with block sparsity.
-        """
         weights = self.conv.weight
         if self.dynamic_pruning:
             return torch.ones_like(weights)
-
         weights_abs = weights.abs()
-        if self.block_size:  # Block sparsity
-            # Reshape weights into blocks and compute block-wise magnitude
+        if self.block_size:
             kh, kw = self.block_size
             weights_reshaped = weights_abs.view(weights_abs.size(0), weights_abs.size(1), 
                                                weights_abs.size(2) // kh, kh, 
@@ -114,7 +111,6 @@ class SparseConv2d(nn.Module):
             block_magnitudes = weights_reshaped.norm(p=2, dim=(3, 4))
             threshold = torch.quantile(block_magnitudes.flatten(), self.sparsity)
             block_mask = (block_magnitudes > threshold).float()
-            # Expand block mask back to full weight shape
             mask = block_mask.unsqueeze(-1).unsqueeze(-1).expand_as(weights_reshaped).reshape_as(weights)
         else:
             threshold = torch.quantile(weights_abs.flatten(), self.sparsity)
@@ -122,7 +118,6 @@ class SparseConv2d(nn.Module):
         return mask
 
     def update_mask(self):
-        """Update mask dynamically based on current weight magnitudes."""
         if self.dynamic_pruning:
             self.mask.data = self.generate_mask()
 
@@ -143,7 +138,7 @@ class SparseConv2d(nn.Module):
 class SparseMLP(nn.Module):
     """
     Sparse MLP with Tensor Core Acceleration and optional dynamic pruning.
-    Uses sparse linear layers to reduce computation.
+    Uses sparse linear layers with fused ops for efficiency.
     """
     def __init__(self, input_dim, hidden_dim, output_dim, sparsity=0.5, 
                  use_fp16=True, dynamic_pruning=False):
@@ -155,9 +150,41 @@ class SparseMLP(nn.Module):
     def forward(self, x):
         if self.use_fp16:
             with autocast():
-                x = F.relu(self.fc1(x))
+                x = self.fc1(x)  # Already includes ReLU from fused kernel
                 x = self.fc2(x)
                 return x
         else:
-            x = F.relu(self.fc1(x.float()))
-            return self.fc2(x.float())
+            x = self.fc1(x)  # Includes ReLU from fallback
+            return self.fc2(x)
+
+# Example training loop with Apex mixed precision and FusedAdam
+def train_sparse_mlp():
+    model = SparseMLP(784, 256, 10, sparsity=0.5, use_fp16=True).cuda()
+    optimizer = FusedAdam(model.parameters(), lr=0.001)
+    
+    # Initialize Apex AMP
+    model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
+
+    # Dummy data
+    inputs = torch.randn(32, 784).cuda()
+    targets = torch.randint(0, 10, (32,)).cuda()
+
+    # Training loop
+    for _ in range(100):
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = F.cross_entropy(outputs, targets)
+        
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+        optimizer.step()
+
+    # Export to ONNX
+    torch.onnx.export(model, inputs, "sparse_mlp.onnx", opset_version=12)
+
+    # Convert to TensorRT
+    model_trt = torch2trt(model, [inputs], fp16_mode=True)
+    return model_trt
+
+if __name__ == "__main__":
+    trt_model = train_sparse_mlp()