(
local_batch,
dp_mesh,
seed,
dataset_args,
num_workers,
bucket_config_type: str = None,
dataset_name: str = 'ZoeDataset',
is_test: bool = False,
)
| 25 | |
| 26 | |
| 27 | def get_dataloader( |
| 28 | local_batch, |
| 29 | dp_mesh, |
| 30 | seed, |
| 31 | dataset_args, |
| 32 | num_workers, |
| 33 | bucket_config_type: str = None, |
| 34 | dataset_name: str = 'ZoeDataset', |
| 35 | is_test: bool = False, |
| 36 | ): |
| 37 | drop_last = False |
| 38 | collate_fn = collate_fn_default |
| 39 | if dataset_name == 'WanMotionTensorDatasetWiRefMotion': |
| 40 | data_cls = WanMotionTensorDatasetWiRefMotion |
| 41 | collate_fn = collate_fn_motion_wanvideo |
| 42 | dataset_args['is_test'] = is_test |
| 43 | elif dataset_name == 'MBenchWiRefMotion': |
| 44 | data_cls = MBenchWiRefMotion |
| 45 | collate_fn = collate_fn_motion_wanvideo |
| 46 | assert is_test == True, "MBench Dataset is for testing only" |
| 47 | else: |
| 48 | data_cls = ZoeMixDataset |
| 49 | dataset_args['dp_rank'] = dp_mesh.get_local_rank() |
| 50 | dataset = data_cls(**dataset_args) |
| 51 | if bucket_config_type is None: |
| 52 | sampler = StatefulDistributedSampler( |
| 53 | dataset, |
| 54 | num_replicas=dp_mesh.size(), |
| 55 | rank=dp_mesh.get_local_rank(), |
| 56 | shuffle=not is_test, |
| 57 | ) |
| 58 | dataloader = DataLoader( |
| 59 | dataset, |
| 60 | batch_size=local_batch, |
| 61 | sampler=sampler, |
| 62 | worker_init_fn=get_seed_worker(seed), |
| 63 | drop_last=drop_last, |
| 64 | pin_memory=True, |
| 65 | num_workers=num_workers, |
| 66 | collate_fn=collate_fn, |
| 67 | prefetch_factor=2, |
| 68 | ) |
| 69 | else: |
| 70 | sampler = DistributedFlopBalanceSampler( |
| 71 | dataset, |
| 72 | dp_rank=dp_mesh.get_local_rank(), |
| 73 | dp_size=dp_mesh.size(), |
| 74 | global_seed=seed, |
| 75 | bucket_config_type=bucket_config_type, |
| 76 | ) |
| 77 | dataloader = DataLoader( |
| 78 | dataset, |
| 79 | batch_size=local_batch, |
| 80 | sampler=sampler, |
| 81 | drop_last=drop_last, |
| 82 | pin_memory=True, |
| 83 | num_workers=num_workers, |
| 84 | ) |
no test coverage detected