How to properly integrate Energon dataloader with Hugging Face Accelerate?

[leo1oel]

Problem Description

I'm trying to use Megatron-Energon's dataloader with Hugging Face Accelerate for distributed training, but encountering a tensor concatenation error when using accelerator.prepare() on the dataloader.

Custom Sample Definition

@dataclass
class CustomSample(Sample):
    cls: torch.Tensor
    image: torch.Tensor
    latent: torch.Tensor
def cook_custom_sample(sample: dict) -> CustomSample:
    return CustomSample(
        **basic_sample_keys(sample),
        cls=torch.tensor(sample["cls"], dtype=torch.long),
        image=sample["image.png"],
        latent=torch.from_numpy(sample["latent.npy"]),
    )

Training Setup

ds = get_train_dataset(
    args.data_dir,
    batch_size=local_batch_size,
    shuffle_buffer_size=100,
    task_encoder=CustomTaskEncoder(),
    max_samples_per_sequence=100,
    worker_config=WorkerConfig.default_worker_config(),
)
train_dataloader = get_savable_loader(ds)
# This line causes the error
model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)

Error Message

TypeError: Can only concatenate tensors but got <class 'int'>

The error occurs in Accelerate's internal concatenation operations when it tries to process the dataloader batches.

Does Energon handle distributed training internally, making Accelerate preparation unnecessary? What's the recommended integration pattern for using both libraries together?

Any guidance on the proper integration pattern would be greatly appreciated!

[philipp-fischer]

Hi @leo1oel, it's great you are trying energon with accelerate. Frankly, we have never done this before so it's no surprise it's not working out of the box.
Accelerate does a few things to the data loader, depending on what type of data loader it finds. Our data loader is quite special and I think we will need some adaptations to make it work with accelerate.
One (untested) approach that might make it work is by setting the energon worker config to single rank and then use dispatch_batches=True with accelerate, so that it only iterates the data on one rank and distributes it to the others.
Note however, that this is not optimal in terms of speed.
For a seamless integration with high speed, some more work on our end will be needed. We'll look into it, but it's not high-prio at the moment.

[leo1oel]

OK, thank you very much!