跳到内容

API

BionemoMegatronModel

基类:MegatronModuleGeneric[DataT]ABC

使用 Megatron 的模型必须是 MegatronModule 类型。

唯一的主要区别是显式的 forward 传递方法签名,这使得此类与 bionemo-core 的 Model 结构类型兼容。

源代码位于 bionemo/llm/api.py
32
33
34
35
36
37
38
39
40
41
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()