Add option for LoRA with Transformer Engine op fuser

.github/workflows/cicd-main-nemo2.yml

@@ -143,6 +143,8 @@ jobs:
             runner: self-hosted-azure
           - script: L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_Chat
             runner: self-hosted-azure
+          - script: L2_NeMo_2_GPT_LoRA_TP1PP1_MBS1_TE_op_fuser
+            runner: self-hosted-azure
           - script: L2_NeMo_2_Mixtral_LoRA_EP2PP1_MBS2_exclude
             runner: self-hosted-azure
           - script: L2_NeMo_2_Mixtral_LoRA_EP2PP1_MBS2

nemo/collections/llm/fn/base.py

@@ -20,6 +20,8 @@
 
 @runtime_checkable
 class HasBool(Protocol):
+    """Protocol for objects with bool operation"""
+
     def __bool__(self) -> bool: ...
 
 
@@ -72,13 +74,18 @@ def double_weights(m):
     if not kwargs.pop("_skip_map", False) and hasattr(module, "map"):
         return module.map(func, leaf_only=leaf_only, **kwargs)
 
+    if isinstance(module, nn.Module) and not isinstance(module, (nn.Sequential, nn.ModuleList, nn.ModuleDict)):
+        return _map_module(module, func, leaf_only=leaf_only, **kwargs)
     elif isinstance(module, Iterable):
+        # Assume iterable is API-compatible with dict or list
         if all(hasattr(module, key) for key in ["items", "values", "keys"]):
             return _map_module_dict(module, func, leaf_only=leaf_only, **kwargs)
-
         return _map_module_list(module, func, leaf_only=leaf_only, **kwargs)
     else:
-        return _map_module(module, func, leaf_only=leaf_only, **kwargs)
+        raise ValueError(
+            "Expected `module` to be a PyTorch module or a collection of modules, "
+            f"but got {module.__class__.__name__}."
+        )
 
 
 def walk(