flow numerical tests + add option to force use fp32 in backward

src/pyjuice/layer/sum_layer.py

@@ -266,7 +266,8 @@ def backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
                  params: torch.Tensor, param_flows: Optional[torch.Tensor] = None,
                  allow_modify_flows: bool = False, propagation_alg: str = "LL", 
                  logspace_flows: bool = False, negate_pflows: bool = False, 
-                 accumulate_ch_flows: bool = False, allow_neg_flows: bool = False, **kwargs) -> None:
+                 accumulate_ch_flows: bool = False, allow_neg_flows: bool = False,
+                 force_use_fp32: bool = False, **kwargs) -> None:
         """
         Computes the forward pass of a sum layer:
         ```
@@ -329,6 +330,7 @@ def backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
                     negate_pflows = negate_pflows, 
                     accumulate_ch_flows = accumulate_ch_flows,
                     allow_neg_flows = allow_neg_flows,
+                    force_use_fp32 = force_use_fp32,
                     **kwargs
                 )
 
@@ -353,6 +355,7 @@ def backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
                     negate_pflows = negate_pflows, 
                     accumulate_ch_flows = accumulate_ch_flows, 
                     allow_neg_flows = allow_neg_flows,
+                    force_use_fp32 = force_use_fp32,
                     **kwargs
                 )
 
@@ -374,6 +377,7 @@ def backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
                     logspace_flows = logspace_flows,
                     negate_pflows = negate_pflows, 
                     allow_neg_flows = allow_neg_flows, 
+                    force_use_fp32 = force_use_fp32,
                     **kwargs
                 )
 
@@ -1267,7 +1271,8 @@ def _backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
                   logspace_flows: bool = False, 
                   negate_pflows: bool = False, 
                   accumulate_ch_flows: bool = False, 
-                  allow_neg_flows: bool = False, **kwargs) -> None:
+                  allow_neg_flows: bool = False, 
+                  force_use_fp32: bool = False, **kwargs) -> None:
         """
         Back pass of sum layers.
 
@@ -1314,7 +1319,7 @@ def _backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
                 partition_id = partition_id, allow_modify_flows = allow_modify_flows,
                 propagation_alg = propagation_alg, logspace_flows = logspace_flows, 
                 negate_pflows = negate_pflows, accumulate_ch_flows = accumulate_ch_flows, 
-                allow_neg_flows = allow_neg_flows, **kwargs
+                allow_neg_flows = allow_neg_flows, force_use_fp32 = force_use_fp32, **kwargs
             )
 
         elif mode == self.SPARSE:
@@ -1514,7 +1519,7 @@ def _backward_block_sparse(self, node_flows: torch.Tensor, element_flows: torch.
                                cs_block_size: int, local_ids: Optional[torch.Tensor] = None,
                                partition_id: int = -1, allow_modify_flows: bool = False, propagation_alg: str = "LL", 
                                logspace_flows: bool = False, negate_pflows: bool = False, accumulate_ch_flows: bool = False, 
-                               allow_neg_flows: bool = False, **kwargs) -> None:
+                               allow_neg_flows: bool = False, force_use_fp32 = False, **kwargs) -> None:
         """
         Back pass of sum layers with block-sparse processing kernel.
 
@@ -1540,7 +1545,8 @@ def _backward_block_sparse(self, node_flows: torch.Tensor, element_flows: torch.
                 propagation_alg = propagation_alg, 
                 logspace_flows = logspace_flows, 
                 accumulate_ch_flows = accumulate_ch_flows, 
-                allow_neg_flows = allow_neg_flows, **kwargs
+                allow_neg_flows = allow_neg_flows, 
+                force_use_fp32 = force_use_fp32, **kwargs
             )
 
         # Flows w.r.t. parameters
@@ -1695,7 +1701,7 @@ def _bk_triton_block_sparse_ele_kernel(node_flows, element_flows, node_mars, ele
                     acc = tl.where(partial_flows_max == -float("inf"),
                         acc,
                         tl.where(partial_flows_max > acc,
-                            tl.log(partial_flows + tl.exp(acc - partial_flows_max) + 1e-24) + partial_flows_max,
+                            tl.log(partial_flows + tl.exp(acc - partial_flows_max) + 1e-32) + partial_flows_max,
                             tl.log(tl.exp(partial_flows_max - acc) * partial_flows + 1.0) + acc
                         )
                     )
@@ -1854,7 +1860,7 @@ def _bk_triton_block_sparse_ele_csmm2_kernel(node_flows, element_flows, node_mar
                     acc = tl.where(partial_flows_max == -float("inf"),
                         acc,
                         tl.where(partial_flows_max > acc,
-                            tl.log(partial_flows + tl.exp(acc - partial_flows_max) + 1e-24) + partial_flows_max,
+                            tl.log(partial_flows + tl.exp(acc - partial_flows_max) + 1e-32) + partial_flows_max,
                             tl.log(tl.exp(partial_flows_max - acc) * partial_flows + 1.0) + acc
                         )
                     )
@@ -1885,7 +1891,8 @@ def _backward_block_sparse_ele_flows(self, node_flows: torch.Tensor, element_flo
                                          parpids: torch.Tensor, cs_block_size: int, local_ids: Optional[torch.Tensor] = None,
                                          partition_id: int = -1, allow_modify_flows: bool = False,
                                          propagation_alg: str = "LL", logspace_flows: bool = False, 
-                                         accumulate_ch_flows: bool = False, allow_neg_flows: bool = False, **kwargs) -> None:
+                                         accumulate_ch_flows: bool = False, allow_neg_flows: bool = False, 
+                                         force_use_fp32: bool = False, **kwargs) -> None:
 
         assert params.dim() == 1, "Expecting a 1D `params`."
 
@@ -1956,7 +1963,7 @@ def _backward_block_sparse_ele_flows(self, node_flows: torch.Tensor, element_flo
         BLOCK_SIZE_K = self.block_size
         allow_modify_flows = 1 if allow_modify_flows else 0
 
-        if TILE_SIZE_M >= 16 and TILE_SIZE_K >= 16 and BLOCK_B >= 16:
+        if TILE_SIZE_M >= 16 and TILE_SIZE_K >= 16 and BLOCK_B >= 16 and not force_use_fp32:
             TL_DOT = 1
         else:
             TL_DOT = 0

src/pyjuice/model/tensorcircuit.py

@@ -354,6 +354,7 @@ def backward(self, inputs: Optional[torch.Tensor] = None,
                  negate_pflows: bool = False,
                  _inner_layers_only: bool = False,
                  _disable_buffer_init: bool = False,
+                 force_use_fp32: bool = False,
                  **kwargs):
         """
         Backward evaluation of the PC that computes node flows as well as parameter flows.
@@ -372,6 +373,7 @@ def backward(self, inputs: Optional[torch.Tensor] = None,
         self._run_params["propagation_alg"] = propagation_alg
         self._run_params["logspace_flows"] = logspace_flows
         self._run_params["negate_pflows"] = negate_pflows
+        self._run_params["force_use_fp32"] = force_use_fp32
 
         assert self.node_mars is not None and self.element_mars is not None, "Should run forward path first."
         if input_layer_fn is None:
@@ -456,7 +458,7 @@ def _run_inner_layers():
                                              param_flows = self.param_flows if compute_param_flows else None,
                                              allow_modify_flows = allow_modify_flows, 
                                              propagation_alg = propagation_alg if isinstance(propagation_alg, str) else propagation_alg[layer_id], 
-                                             logspace_flows = logspace_flows, negate_pflows = negate_pflows, **kwargs)
+                                             logspace_flows = logspace_flows, negate_pflows = negate_pflows, force_use_fp32 = force_use_fp32, **kwargs)
 
                         # Execute post-backward callback
                         layer_group.callback(
@@ -723,7 +725,8 @@ def get_node_flows(self, ns: CircuitNodes, **kwargs):
                 param_flows = None, allow_modify_flows = self._run_params["allow_modify_flows"], 
                 propagation_alg = self._run_params["propagation_alg"], 
                 logspace_flows = self._run_params["logspace_flows"], 
-                negate_pflows = self._run_params["negate_pflows"], **kwargs
+                negate_pflows = self._run_params["negate_pflows"],
+                force_use_fp32 = self._run_params["force_use_fp32"], **kwargs
             )
 
             return self.element_flows[nsid:neid,:].detach()

tests/model/flow_numerical_test.py

@@ -0,0 +1,115 @@
+import torch
+
+import pyjuice as juice
+
+import pytest
+
+
+def test_hmm_flow_stability():
+
+    device = torch.device("cuda:0")
+
+    seq_length = 16
+    num_latents = 1024
+    vocab_size = 1023
+    batch_size = 32
+
+    for num_latents in [256, 512, 1024, 2048]:
+
+        root_ns = juice.structures.HMM(
+            seq_length = seq_length,
+            num_latents = num_latents,
+            num_emits = vocab_size
+        )
+        root_ns.init_parameters(perturbation = 32.0)
+
+        pc = juice.compile(root_ns)
+        pc.to(device)
+
+        data = torch.randint(0, vocab_size, [batch_size, seq_length], device = device)
+
+        lls = pc(data)
+        pc.backward(data, allow_modify_flows = False, logspace_flows = True)
+
+        pc.update_parameters()
+        pc.update_param_flows()
+
+        alpha = root_ns.chs[0].chs[1].get_source_ns().get_params()
+        block_size = alpha.size(1)
+        num_node_blocks = num_latents // block_size
+        alpha = alpha.reshape(num_node_blocks, num_node_blocks, block_size, block_size).permute(
+            0, 2, 1, 3
+        ).reshape(num_latents, num_latents).to(device)
+
+        alpha_flow = root_ns.chs[0].chs[1].get_source_ns().get_param_flows().reshape(
+            num_node_blocks, num_node_blocks, block_size, block_size).permute(
+            0, 2, 1, 3
+        ).reshape(num_latents, num_latents).to(device)
+
+        gamma = root_ns.get_params().reshape(num_latents).to(device)
+
+        # Root node
+        pc_nflows = pc.get_node_flows(root_ns)
+        assert torch.all(pc_nflows.abs() < 1e-4)
+
+        # First product layer
+        pc_cflows = pc.get_node_flows(root_ns.chs[0])
+        pc_cflows_chk = pc.get_node_flows(root_ns.chs[0].chs[0])
+        assert torch.all(pc_cflows.logsumexp(dim = 0) < 1e-4)
+        assert torch.all((pc_cflows - pc_cflows_chk).abs() < 1e-4)
+
+        nmars = pc.get_node_mars(root_ns)
+        cmars = pc.get_node_mars(root_ns.chs[0])
+        cflows = gamma[:,None] * (cmars - nmars).exp()
+        assert torch.all((pc_cflows - cflows.log()).abs() < 1e-4)
+
+        # The following product layers
+        nflows = cflows.log()
+        ns = root_ns.chs[0].chs[1]
+        layer_id = len(pc.inner_layer_groups) - 3
+        for _ in range(seq_length - 2):
+            cs = ns.chs[0]
+
+            pc_cflows = pc.get_node_flows(cs)
+            pc_cflows_chk = pc.get_node_flows(cs.chs[0])
+            assert torch.all((pc_cflows - pc_cflows_chk).abs() < 1e-4)
+
+            nmars = pc.get_node_mars(ns)
+            cmars = pc.get_node_mars(cs)
+
+            xmatr = alpha[:,:,None] * (cmars[None,:,:] - nmars[:,None,:]).exp()
+            cflows = (nflows[:,None,:] + xmatr.log()).logsumexp(dim = 0)
+
+            # Test node flows
+            assert torch.all(((pc_cflows.exp() - cflows.exp()).abs() / cflows.exp()) < 1e-2)
+            assert torch.all((pc_cflows.exp() - cflows.exp()).abs().mean() < 1e-4)
+
+            pc.zero_param_flows()
+
+            prod_layer = pc.inner_layer_groups[layer_id-1][0]
+            prod_layer.forward(pc.node_mars, pc.element_mars, _for_backward = True)
+            layer = pc.inner_layer_groups[layer_id][0]
+            layer.backward(pc.node_flows, pc.element_flows, pc.node_mars, pc.element_mars,
+                        pc.params, pc.param_flows, logspace_flows = True)
+
+            pc.update_param_flows()
+
+            local_alpha_flow = ns.get_source_ns().get_param_flows().reshape(
+                num_node_blocks, num_node_blocks, block_size, block_size).permute(
+                0, 2, 1, 3
+            ).reshape(num_latents, num_latents).to(device)
+
+            xmatr = alpha[:,:,None] * (cmars[None,:,:] - nmars[:,None,:]).exp()
+            pflows = (nflows[:,None,:] + xmatr.log()).logsumexp(dim = 2)
+
+            # Test parameter flows
+            torch.all((local_alpha_flow - pflows.exp()).abs() / batch_size < 4e-5)
+            torch.all(((local_alpha_flow.sum(dim = 1) - batch_size) / batch_size).abs() < 1e-6)
+
+            ns = ns.chs[0].chs[1]
+            nflows = cflows
+            layer_id -= 2
+
+
+if __name__ == "__main__":
+    test_hmm_flow_stability()

tests/structures/hmm_correctness_test.py

@@ -6,6 +6,8 @@
 import pyjuice as juice
 import pyjuice.nodes.distributions as dists
 
+import pytest
+
 
 def test_hmm_forward_backward():
 
@@ -15,8 +17,8 @@ def test_hmm_forward_backward():
     vocab_size = 1023
     batch_size = 32
 
-    num_node_blocks = 4 # 4096 // 32 # 4
-    block_size = 1024 # 32 # 1024
+    num_node_blocks = 1 # 4096 // 32 # 4
+    block_size = 4096 # 32 # 1024
     num_latents = block_size * num_node_blocks
 
     with juice.set_block_size(block_size = block_size):