Fix DDP parameter and gradient layouts

backbones.py

@@ -151,6 +151,32 @@ class AgeSinusoidalEncoding(nn.Module):
         return output
 
 
+class DepthwiseConv2d(nn.Module):
+    """Depthwise convolution without a singleton weight dimension."""
+
+    def __init__(self, channels: int, kernel_size: int):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.padding = kernel_size // 2
+        self.weight = nn.Parameter(
+            torch.empty(channels, kernel_size, kernel_size)
+        )
+        self.bias = nn.Parameter(torch.empty(channels))
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        bound = 1 / math.sqrt(kernel_size * kernel_size)
+        nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.conv2d(
+            x,
+            self.weight.unsqueeze(1),
+            self.bias,
+            padding=self.padding,
+            groups=self.channels,
+        )
+
+
 class LiteTimesBackbone2d(nn.Module):
     """Cheap local feature extractor for a TimesNet period image."""
 
@@ -162,10 +188,7 @@ class LiteTimesBackbone2d(nn.Module):
         if expansion <= 0:
             raise ValueError("expansion must be > 0")
         hidden_dim = max(dim, int(round(dim * expansion)))
-        self.dwconv = nn.Conv2d(
+        self.dwconv = DepthwiseConv2d(dim, kernel_size)
-            dim, dim, kernel_size=kernel_size,
-            padding=kernel_size // 2, groups=dim,
-        )
         self.norm = nn.GroupNorm(1, dim)
         self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
         self.act = nn.GELU()

models.py

@@ -100,8 +100,11 @@ class DeepHealth(nn.Module):
         self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
         if target_mode == "next_token":
             self.risk_head.weight = self.token_embedding.weight
-        self.query_token = nn.Parameter(torch.zeros(n_embd))
+        if target_mode == "all_future":
-        nn.init.normal_(self.query_token, mean=0.0, std=0.02)
+            self.query_token = nn.Parameter(torch.zeros(n_embd))
+            nn.init.normal_(self.query_token, mean=0.0, std=0.02)
+        else:
+            self.register_parameter("query_token", None)
 
     def _make_history_attn_mask(
         self,
@@ -251,6 +254,11 @@ class DeepHealth(nn.Module):
         h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 
         if mode == "all_future":
+            if self.query_token is None:
+                raise RuntimeError(
+                    "all_future forward requires a model initialized with "
+                    "target_mode='all_future'"
+                )
             batch_size = event_seq.size(0)
             query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
             h_disease = torch.cat([h_disease, query], dim=1)

train_next_step.py

@@ -901,4 +901,8 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    finally:
+        if dist.is_initialized():
+            dist.destroy_process_group()