src/megatron/bridge/models/falcon_h1/modeling_falconh1/falconh1_layer.py (2)

1-1: Update copyright header.

Add the complete Apache 2.0 license text. The copyright year can remain 2024 or be updated to 2025.

---

3-21: Remove unused imports.

Static analysis indicates several unused imports that should be removed for code cleanliness.

🧹 Proposed cleanup

-import math
 from dataclasses import dataclass
 from typing import Optional, Tuple, Union

 import torch
-import torch.nn as nn
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.bridge.models.falcon_h1.modeling_falconh1.falconh1_model import FalconH1Config
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec, build_module
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.utils import log_single_rank
 import logging
-from megatron.core.ssm.mamba_mixer import MambaMixer, MambaMixerSubmodules
+from megatron.core.ssm.mamba_mixer import MambaMixerSubmodules
-from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules
+from megatron.core.transformer.attention import SelfAttentionSubmodules
 from megatron.core.transformer.mlp import MLP

src/megatron/bridge/models/falcon_h1/modeling_falconh1/falconh1_model.py (3)

1-1: Update copyright header to 2025 with full license text.

---

20-21: Remove duplicate Optional and Tuple imports.

These types are already imported on line 3. The redefinition is flagged by static analysis.

🧹 Proposed fix

-from dataclasses import dataclass
-from typing import Optional, Tuple
+from dataclasses import dataclass

The Tuple type from line 3 can be used, and Optional is already imported.

---

177-177: Use explicit str | None for nullable type hints.

Per coding guidelines, use T | None instead of implicit Optional. PEP 484 prohibits implicit Optional.

📝 Proposed fix

-        hybrid_override_pattern: str = None,
+        hybrid_override_pattern: str | None = None,

src/megatron/bridge/models/falcon_h1/modeling_falconh1/falconh1_block.py (2)

402-406: Remove unused output variable.

The make_viewless_tensor result is assigned to output but hidden_states is returned instead. Either use output or remove the assignment.

🧹 Proposed fix

         # Ensure that the tensor passed between pipeline parallel stages is
         # viewless. See related notes in TransformerBlock and TransformerLayer
-        output = make_viewless_tensor(
+        hidden_states = make_viewless_tensor(
             inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True
         )

         return hidden_states

---

131-131: Minor style issues flagged by static analysis.

1. Line 131: Use explicit str | None for nullable type hint
2. Line 453: Use is not instead of not ... is

📝 Proposed fixes

-        hybrid_override_pattern: str = None,
+        hybrid_override_pattern: str | None = None,

-            if not module is self.layers:
+            if module is not self.layers:

Also applies to: 453-453

src/megatron/bridge/models/falcon_h1/modeling_falconh1/mamba_hybrid_layer_allocation.py (6)

1-1: Update copyright year to 2025.

The copyright header uses 2024, but per the coding guidelines and consistency with other files in this PR (e.g., falconh1_provider.py), it should be 2025.

-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

---

17-22: Consider using frozenset for immutability or annotate with ClassVar.

The VALID set is a mutable class attribute. For a constant set that shouldn't be modified, consider using frozenset for true immutability, or at minimum annotate with ClassVar to satisfy the linter.

+from typing import ClassVar
+
 class Symbols:
     MAMBA = "M"
     ATTENTION = "*"
     MLP = "-"
     PARALLEL = "P"
-    VALID = {MAMBA, ATTENTION, MLP, PARALLEL}
+    VALID: ClassVar[frozenset[str]] = frozenset({MAMBA, ATTENTION, MLP, PARALLEL})

---

107-108: Use explicit str | None type hint.

Per coding guidelines, use T | None for nullable types instead of implicit Optional. The current signature uses str = None which implicitly means Optional[str].

-    override_pattern: str = None,
+    override_pattern: str | None = None,

As per coding guidelines: "Use 'T | None' for nullable types instead of 'Optional[T]'"

---

109-113: Consider splitting compound assertions for clearer error messages.

When a compound assertion fails, it's not immediately clear which condition was violated. Splitting them provides better debugging information.

♻️ Proposed fix

     assert total_layers_count > 0
-    assert target_attention_ratio >= 0.0 and target_attention_ratio <= 1.0
-    assert target_mlp_ratio >= 0.0 and target_mlp_ratio <= 1.0
-    assert target_parallel_hybrid_ratio >= 0.0 and target_parallel_hybrid_ratio <= 1.0
+    assert 0.0 <= target_attention_ratio <= 1.0, f"target_attention_ratio must be in [0, 1], got {target_attention_ratio}"
+    assert 0.0 <= target_mlp_ratio <= 1.0, f"target_mlp_ratio must be in [0, 1], got {target_mlp_ratio}"
+    assert 0.0 <= target_parallel_hybrid_ratio <= 1.0, f"target_parallel_hybrid_ratio must be in [0, 1], got {target_parallel_hybrid_ratio}"
     assert target_attention_ratio + target_mlp_ratio + target_parallel_hybrid_ratio <= 1.0

---

25-27: Add parameterized return type hints.

The return type list could be more specific as list[str] to improve type safety and IDE support. This applies to _allocate_auto, _allocate_override, and allocate_layers.

 def _allocate_auto(
     total_layers_count: int, target_attention_ratio: float, target_mlp_ratio: float, target_parallel_hybrid_ratio: float
-) -> list:
+) -> list[str]:

---

133-135: Use logging.WARNING level for warning messages.

The message says "Warning:" but uses logging.INFO level. Consider using logging.WARNING for consistency with the message semantics.

-            log_single_rank(logger, logging.INFO, "Warning: overriding pattern A with pattern B")
+            log_single_rank(logger, logging.WARNING, "Overriding pattern A with pattern B")

src/megatron/bridge/models/falcon_h1/falconh1_provider.py (3)

17-17: Prefer X | None and X | Y syntax over Optional and Union.

Per coding guidelines, use T | None for nullable types and X | Y for union types instead of Optional[T] and Union[X, Y]. The Optional and Union imports can be removed if you update the type hints.

Lines affected:

• Line 67: Optional[str] → str | None
• Line 73: Optional[float] → float | None
• Line 78: Optional[int] → int | None
• Line 107: Union[ModuleSpec, Callable[[], ModuleSpec]] → ModuleSpec | Callable[[], ModuleSpec]

As per coding guidelines: "Use 'T | None' for nullable types instead of 'Optional[T]'" and "Use 'X | Y' for union types instead of 'Union[X, Y]'"

---

94-94: Fix typo in comment.

-    `#Falcon` H1 Mup Fwd Multpliers
+    # Falcon H1 MuP Forward Multipliers

---

102-103: Add parameterized tuple type hints.

For better type safety, specify the tuple element types:

-    mlp_multipliers: tuple = (1.0, 1.0)
-    ssm_multipliers: tuple = (1.0, 1.0, 1.0, 0.5, 1.0)
+    mlp_multipliers: tuple[float, float] = (1.0, 1.0)
+    ssm_multipliers: tuple[float, float, float, float, float] = (1.0, 1.0, 1.0, 0.5, 1.0)