Energon 多模态数据加载框架

from megatron.energon import get_train_dataset, get_loader, WorkerConfig

# 创建数据集
ds = get_train_dataset(
    '/my/dataset/path',
    batch_size=1,
    shuffle_buffer_size=100,
    max_samples_per_sequence=100,
    worker_config=WorkerConfig.default_worker_config(),
)

# 创建数据加载器
loader = get_loader(ds)

# 迭代训练数据
for batch in loader:
    # 执行推理、梯度更新等操作
    pass

CaptioningSample(
    __key__='part_00123/00403.tar/004030195',
    image=tensor(...),  # 形状: (1, 3, H, W)
    caption='...'
)

loader = get_loader(get_train_dataset(
    '/path/to/your/dataset',
    batch_size=2,
    shuffle_buffer_size=100,
    max_samples_per_sequence=100,
    worker_config=WorkerConfig.default_worker_config(),
))

for batch in loader:
    print(batch)
    break

from megatron.energon import get_train_dataset, get_savable_loader, WorkerConfig

ds = get_train_dataset(..., worker_config=wc)
loader = get_savable_loader(ds)

# 训练循环中
for step, batch in enumerate(loader):
    # 执行训练步骤
    if should_checkpoint(step):
        state = loader.save_state_rank()  # 与当前 DP rank 绑定
        # 将状态保存到 checkpoint
        # torch.save(state, path)

ds = get_train_dataset(..., worker_config=wc)  # worker_config 必须与保存时一致
loader = get_savable_loader(ds)

# 从 checkpoint 加载状态
state = torch.load(path)['loader_states'][dp_rank]
loader.restore_state_rank(state)

# 继续迭代,从中断处的下一条样本开始
for batch in loader:
    # 训练逻辑
    pass

class MyTaskEncoder(DefaultTaskEncoder):
    decoder = SampleDecoder(image_decode="pilrgb")

ds = get_train_dataset(
    path,
    batch_size=4,
    auto_decode=True,  # 或 False
)

decoder = SampleDecoder(av_decode="AVDecoder")

# 提取前 2 秒和第 4-6 秒的片段
dat = av_decoder.get_clips(
    video_clip_ranges=[(0, 2), (4, 6)],
    video_unit="seconds",
)

print(dat.video_clips[0].shape)  # 第一个片段
print(dat.video_clips[1].shape)  # 第二个片段

from megatron.energon.av import get_clips_uniform

dat = get_clips_uniform(
    av_decoder=av_decoder,
    num_clips=5,              # 提取 5 个片段
    clip_duration_seconds=3,  # 每个片段 3 秒
    request_audio=True        # 同时返回音频
)

frames = av_decoder.get_single_frames_uniform()

dat = av_decoder.get_video()              # 仅视频
dat = av_decoder.get_audio()              # 仅音频
dat = av_decoder.get_video_with_audio()   # 视频+音频

def cook_crude_sample(self, sample: Union[T_sample, CrudeSample]) -> T_sample

def encode_sample(self, sample: T_sample) -> T_encoded_sample

def preencode_sample(self, sample: T_sample) -> T_sample

    def select_samples_to_pack(self, samples: List[CaptioningSample]) -> List[List[CaptioningSample]]:
        # Do something intelligent here, e.g. sort by caption length and concat where possible.
        # This could be better, but it's just an example.
        samples.sort(key=lambda x: len(x.caption))
        groups = []
        while len(samples) > 0:
            batch = []
            caption_len = 0
            while len(samples) > 0 and caption_len + len(samples[0].caption) < self.max_length:
                sample = samples.pop(0)
                batch.append(sample)
                caption_len += len(sample.caption)
            groups.append(batch)
        return groups
    

ds = get_train_dataset(
    ...,
    packing_buffer_size=1000,  # buffer 越大,越容易找到合适的组合
)

def postencode_sample(self, sample: T_sample) -> T_encoded_sample

def pack_selected_samples(self, samples: List[T_encoded_sample]) -> T_batch_sample

def batch(self, batch: List[T_encoded_sample]) -> T_raw_batch

    @stateless
    def encode_batch(self, batch: T_raw_batch) -> Union[T_batch, Generator[T_batch, None, None]]:
        """Encode a batch of samples. The default implementation converts to the
        _encoded_batch_type."""
        if self._batch_type is None or self._raw_batch_type == self._batch_type:
            return batch
        if is_dataclass(batch):
            fields = {field.name: getattr(batch, field.name) for field in dataclasses.fields(batch)}
        elif isinstance(batch, tuple) and hasattr(batch, "_fields"):
            fields = {field: getattr(batch, field) for field in batch._fields}
        elif isinstance(batch, dict):
            fields = batch
        else:
            raise ValueError("Unrecognized sample type.")
        if issubclass(self._batch_type, dict):
            return fields
        elif dataclasses.is_dataclass(self._batch_type) or issubclass(self._batch_type, tuple):
            # DataClass or NamedTuple
            return self._batch_type(**fields)
        else:
            raise ValueError("Unrecognized encoded sample type.")

for batch in loader:
    batch = batch.pin_memory()  # 或 DataLoader 自动调用
    batch = batch.to(device, non_blocking=True)

from dataclasses import dataclass
from typing import Callable, List, Optional

import torch
from transformers import PreTrainedTokenizerBase as Tokenizer

from megatron.energon import (
    Batch, CaptioningSample, DefaultTaskEncoder,
    SampleDecoder
)

# 中间态批(collate 后、encode_batch 前)
@dataclass
class CaptioningRawBatch(Batch):
    image: torch.Tensor      # (n, c, h, w)
    caption: List[str]       # (n,)

# 最终送入模型的批
@dataclass
class CaptioningBatch(Batch):
    images: torch.Tensor         # (n, c, h, w)
    text_tokens: torch.Tensor    # (n, L)
    text_attn_mask: torch.Tensor # (n, L)

class CaptioningTaskEncoder(
    DefaultTaskEncoder[
        CaptioningSample,
        CaptioningSample,
        CaptioningRawBatch,
        CaptioningBatch
    ]
):
    """图像描述任务编码器"""

    # 声明解码器: 图像→RGB torch.Tensor
    decoder = SampleDecoder(image_decode="torchrgb")

    def __init__(
        self,
        tokenizer: Tokenizer,
        image_transform: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        max_length: int = 128,
    ):
        super().__init__(batch_type=CaptioningRawBatch)
        self.tokenizer = tokenizer
        self.image_transform = image_transform
        self.max_length = max_length

    def encode_sample(self, sample: CaptioningSample) -> CaptioningSample:
        """单样本级处理: 图像增强/变换"""
        if self.image_transform is not None:
            sample.image = self.image_transform(sample.image)
        return sample
    
    def batch(self, samples: List[CaptioningSample]) -> CaptioningRawBatch:
        """批处理"""
        return CaptioningRawBatch.from_samples(samples)

    def encode_batch(self, batch_data: CaptioningRawBatch) -> CaptioningBatch:
        """整批编码: 一次性 tokenize"""
        tokenized = self.tokenizer(
            batch_data.caption,
            padding=True,
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt",
        )
        return CaptioningBatch.derive_from(
            batch_data,
            images=batch_data.image,
            text_tokens=tokenized["input_ids"],
            text_attn_mask=tokenized["attention_mask"],
        )

import torch
from megatron.energon import get_train_dataset, get_savable_loader, WorkerConfig

# 初始化进程组
torch.distributed.init_process_group(backend='nccl')

# 获取各维度的 rank
global_rank = torch.distributed.get_rank()
data_parallel_rank = [0, 0, None, None, 1, 1, None, None][global_rank]
pipeline_parallel_rank = [0, 0, 1, 1, 0, 0, 1, 1][global_rank]
tensor_parallel_rank = [0, 1, 0, 1, 0, 1, 0, 1][global_rank]

# 创建数据并行组
if global_rank in (0, 4):
    data_parallel_group = torch.distributed.new_group(ranks=[0, 4])
elif global_rank in (1, 5):
    data_parallel_group = torch.distributed.new_group(ranks=[1, 5])
else:
    data_parallel_group = None

if data_parallel_rank is not None:
    assert pipeline_parallel_rank == 0, "仅流水线并行 rank 0 加载数据"
    
    # 配置 worker
    worker_config = WorkerConfig(
        rank=data_parallel_rank,
        world_size=torch.distributed.get_world_size(data_parallel_group),
        num_workers=3,
        data_parallel_group=data_parallel_group,
    )

    # 创建加载器
    loader = get_savable_loader(get_train_dataset(
        'coyo-coco-dataset.yaml',
        batch_size=4,
        shuffle_buffer_size=100,
        max_samples_per_sequence=100,
        worker_config=worker_config,
    ))

    # 训练循环
    for i, batch in zip(range(10), loader):
        print(batch)
        break

    # 保存状态(仅第一个 TP rank)
    if tensor_parallel_rank == 0:
        state = loader.save_state_rank()
        torch.save(state, f"dataloader_rank{data_parallel_rank}.pt")

# 恢复状态
if data_parallel_rank is not None:
    assert pipeline_parallel_rank == 0, "仅流水线并行 rank 0 加载数据"

    # 创建新加载器
    loader = get_savable_loader(get_train_dataset(
        'coyo-coco-dataset.yaml',
        batch_size=4,
        shuffle_buffer_size=100,
        max_samples_per_sequence=100,
        worker_config=worker_config,
    ))

    # 加载并恢复状态
    state = torch.load(f"dataloader_rank{data_parallel_rank}.pt")
    loader.restore_state_rank(state)

    # 继续训练
    for batch in loader:
        # 训练逻辑
        pass