Question about len(StreamingDataLoader) on multiple nodes

[Aceticia]

Hi, I've been training a model on 4x8GPU nodes with fabric on a local slurm machine. It has worked well but a bit slow when loading data, so I'm trying to use litdata.optimize and streaming dataset to speed up the loading. I find that when I change the number of nodes provided to the process, fabric does recognize the number of nodes and prints out the correct rank but len(StreamingDataLoader) does not change regardless of the number of nodes. What could be happening here?

To be more specific, the dataset itself is length N, and len(dataloader) also just happens to be N // batch_size instead of N // (total device count x batch_size). It also doesn't look like this is the global batch size, since the memory usage per device is consistent when I change the number of nodes.

Here's how I'm setting up the dataset and fabric:

# Setup dataset
dataset = StreamingDataset(optimized_data_path, shuffle=True, drop_last=True)
train_loader = StreamingDataLoader(
    dataset,
    batch_size=data_config_dict["train_batch_size"],
    num_workers=data_config_dict["num_workers"],
    prefetch_factor=data_config_dict["prefetch_factor"],
    pin_memory=data_config_dict["pin_memory"],
)

# Setup fabric
fabric = L.Fabric(
    precision="bf16-mixed",
    loggers=[logger],
    devices="auto",
    num_nodes=int(os.environ["SLURM_JOB_NUM_NODES"]),
)
fabric.launch()

# Setup dataloader
train_loader = fabric.setup_dataloaders(train_loader)

# Train
for batch in tqdm(train_laoder, total=len(train_loader)):
...