使用 PyTorch 和 NVIDIA BioNeMo Recipes 扩展生物学 Transformer 模型

import torch
import transformer_engine.pytorch as te
from transformer_engine.common.recipe import Format, DelayedScaling

class MyEsmEncoder(torch.nn.Module):
    def __init__(self, num_layers, hidden_size, ffn_hidden_size, num_heads):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            te.TransformerLayer(
                hidden_size=hidden_size,
                ffn_hidden_size=ffn_hidden_size,
                num_attention_heads=num_heads,
                layer_type="encoder",
                self_attn_mask_type="padding",
                attn_input_format="bshd", # or 'thd', read below.
                window_size=(-1, -1), # disable windowed attention
            ) for _ in range(num_layers)
        ])
        # Optionally add embedding, head, etc.

    def forward(self, x, attention_mask=None):
        for layer in self.layers:
            x = layer(x, attention_mask=attention_mask)
        return x

# Layer configuration
layer_num = 8
hidden_size = 4096
sequence_length = 2048
batch_size = 4
ffn_hidden_size = 16384
num_attention_heads = 32
dtype = torch.bfloat16

# Synthetic data (batch, seq, hidden) for bshd format
x = torch.rand(batch_size, sequence_length, hidden_size).cuda().to(dtype=dtype)
attention_mask = torch.ones(batch_size, 1, 1, sequence_length, dtype=torch.bool).cuda()
myEsm = MyEsmEncoder(layer_num, hidden_size, ffn_hidden_size, num_attention_heads)
myEsm.to(dtype=dtype).cuda()

fp8_recipe = DelayedScaling(fp8_format=Format.HYBRID)
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    y = myEsm(x, attention_mask=attention_mask)

def sequence_pack(input_ids, labels):
    # input_ids is a list of sequences: [(S1,), (S2,), ..., (SN,)] of shape (B,S)
    # Flatten and track sequence boundaries

    # Determine the length of each sequence    
    sample_lengths = [len(sample) for sample in input_ids]

    # Flatten the input_ids and labels
    flat_input_ids = [token for sample in input_ids for token in sample]
    flat_labels = [label for sample in labels for label in sample]

    # Create a list of cumulative sums showing where the sequences start/stop
    # Note: for self attention cu_seqlens_q and cu_seqlens_kv will be the same
    cu_seqlens = torch.cumsum(torch.tensor([0] + sample_lengths), dim=0, dtype=torch.int32)

    max_length = max(sample_lengths)
    
    return {
        "input_ids": torch.tensor(flat_input_ids, dtype=torch.int64),
        "labels": torch.tensor(flat_labels, dtype=torch.int64),
        # These are the same kwargs used by `flash_attn_varlen_func`, etc.
        "cu_seqlens_q": cu_seqlens,
        "cu_seqlens_kv": cu_seqlens,
        "max_length_q": max_length,
        "max_length_kv": max_length,
    }