Api

BionemoMegatronModel

Bases: MegatronModule, Generic[DataT], ABC

Models that use Megatron must be a MegatronModule type.

The only major difference is the explicit forward pass method signature that makes this class compatible with bionemo-core's Model structural type.

Source code in bionemo/llm/api.py

class BionemoMegatronModel(MegatronModule, Generic[DataT], ABC):
    """Models that use Megatron must be a MegatronModule type.

    The only major difference is the explicit `forward` pass method signature that makes this class compatible
    with bionemo-core's `Model` structural type.
    """

    @abstractmethod
    def forward(self, *args, **kwargs) -> DataT:  # noqa: D102
        raise NotImplementedError()