Add ctx function for no syn FSDP when wrapping TiledMLP

Author: alektebel

src/liger_kernel/ops/utils.py

@@ -153,12 +153,9 @@ def set_large_grf_mode(kernel_args: dict):
         kernel_args["grf_mode"] = "large"
 
 
-def _get_fsdp_ctx(mlp_module: torch.nn.Module):
-  """
-  Return FSDP.summon_full_params context if module is FSDP-wrapped,
-  otherwise return a no-op context.
-  """
-  if isinstance(mlp_module, FSDP):
-      return FSDP.summon_full_params(mlp_module, write_back=True)
-  else:
-      return contextlib.nullcontext()
+def _get_no_sync_context(module):
+    """Return no_sync context if module is DDP or FSDP, else a no-op."""
+    if isinstance(module, (FSDP, parallel.DistributedDataParallel)):
+        return module.no_sync()
+    import contextlib
+    return contextlib.nullcontext()