thunderfx: Update unsupported ctx with new functorch increment/decrement (#2732)

Author: kshitij12345

thunder/dynamo/utils.py

@@ -399,17 +399,32 @@ def get_nodes_in_unsupported_ctx_regions(gm: torch.fx.GraphModule) -> set[torch.
     # NOTE - Currently doesn't ban any ctx (previously used for `no_grad` and `autocast`).
 
     nodes_in_unsupported_ctx_regions: set[torch.fx.Node] = set()
-    ctx_cnt = 0  # Count of  we have seen till now
+    ctx_stack = []  # Unsupported context.__enter__ we have seen till now.
+
+    CTX_ENTER_EXIT_MAP = {
+        torch._C._functorch._vmap_increment_nesting: torch._C._functorch._vmap_decrement_nesting,
+    }
+
+    # Older version of PyTorch may not have `torch._functorch.predispatch`
+    if hasattr(torch._functorch, "predispatch"):
+        CTX_ENTER_EXIT_MAP[torch._functorch.predispatch._vmap_increment_nesting] = (
+            torch._functorch.predispatch._vmap_decrement_nesting
+        )
+
+    UNSUPPORTED_THUNDER_CTX_ENTER = tuple(CTX_ENTER_EXIT_MAP.keys())
+    UNSUPPORTED_THUNDER_CTX_EXIT = tuple(CTX_ENTER_EXIT_MAP.values())
 
-    UNSUPPORTED_THUNDER_CTX = (torch._C._functorch._vmap_increment_nesting, torch._C._functorch._vmap_decrement_nesting)
     for node in gm.graph.nodes:
-        if node.op == "call_function" and node.target in UNSUPPORTED_THUNDER_CTX:
-            ctx_cnt += 1
-        elif node.op == "call_function" and node.target in UNSUPPORTED_THUNDER_CTX:
-            ctx_cnt -= 1
-        else:
-            if ctx_cnt > 0:
-                nodes_in_unsupported_ctx_regions.add(node)
+        # All the cases when `node` should be marked as in unsupported ctx region.
+        if node.op == "call_function" and node.target in UNSUPPORTED_THUNDER_CTX_ENTER:
+            ctx_stack.append(node.target)
+            nodes_in_unsupported_ctx_regions.add(node)
+        elif node.op == "call_function" and node.target in UNSUPPORTED_THUNDER_CTX_EXIT:
+            enter_fn = ctx_stack.pop()
+            assert CTX_ENTER_EXIT_MAP[enter_fn] == node.target, "Mismatched ctx enter-exit"
+            nodes_in_unsupported_ctx_regions.add(node)
+        elif len(ctx_stack) > 0:
+            nodes_in_unsupported_ctx_regions.add(node)
 
     return nodes_in_unsupported_ctx_regions
 

thunder/tests/test_networks.py

@@ -357,9 +357,6 @@ def test_quantization():
         assert_close(v, sd2[k])
 
 
-@pytest.mark.skip(
-    reason="incompatible with transformers >= 4.55.4, see https://github.com/Lightning-AI/lightning-thunder/issues/2726"
-)
 @thunder.tests.framework.requiresCUDA
 def test_thunderfx_mistral_nemo_small():
     """
@@ -402,7 +399,11 @@ def test_thunderfx_mistral_nemo_small():
     input_ids = torch.randint(0, config.vocab_size, iid_size, device=device)
     attention_mask = torch.ones_like(input_ids)
 
-    output = mdl(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=FutureWarning, message=r".*`isinstance\(treespec, LeafSpec\)` is deprecated.*"
+        )
+        output = mdl(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
     logits = output.logits
     grad_logits = torch.randn_like(logits)
     logits.backward(grad_logits)
@@ -424,9 +425,6 @@ def qwen2():
     return [(phi3), (qwen2)]
 
 
-@pytest.mark.skip(
-    reason="incompatible with transformers >= 4.55.4, see https://github.com/Lightning-AI/lightning-thunder/issues/2726"
-)
 @thunder.tests.framework.requiresCUDA
 @pytest.mark.parametrize("model_fn", _get_model_config_pairs())
 def test_hf_for_nemo(model_fn):
@@ -460,7 +458,11 @@ def test_hf_for_nemo(model_fn):
     ref_output = model(input_ids=input_ids, labels=input_ids)
     ref_loss = ref_output.loss
 
-    compiled_output = compiled_model(input_ids=input_ids, labels=input_ids)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=FutureWarning, message=r".*`isinstance\(treespec, LeafSpec\)` is deprecated.*"
+        )
+        compiled_output = compiled_model(input_ids=input_ids, labels=input_ids)
     compiled_loss = compiled_output.loss
 
     # Less strict tolerance probably due to different type promotion order for bfloat16
@@ -523,6 +525,7 @@ def test_hf_for_nemo(model_fn):
 @requiresCUDA
 @pytest.mark.skip(
     reason="incompatible with transformers >= 4.55.4, see https://github.com/Lightning-AI/lightning-thunder/issues/2726"
+    "Error Message: 'DynamicCache' object has no attribute 'get_usable_length'. Did you mean: 'get_seq_length'?"
 )
 @requiresDeviceMemory(required_memory_bytes=int(0.7 * 1024 * 1024 * 1024))
 @pytest.mark.parametrize(