Spaces:
Running on Zero
Running on Zero
File size: 3,030 Bytes
fed6c68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | # Copyright 2025 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING, Any, Dict, Optional
from torch.utils.data import IterableDataset
from torchdata.stateful_dataloader.sampler import StatefulDistributedSampler
from veomni.utils.device import get_device_type
from ...distributed.parallel_state import get_parallel_state
from ...utils import logging
from ..data_collator import MainCollator, MakeMicroBatchCollator, NoopDataCollator
from ..data_loader import DistributedDataloader
if TYPE_CHECKING:
from torch.utils.data import Dataset
logger = logging.get_logger(__name__)
def build_dit_dataloader(
dataset: "Dataset",
micro_batch_size: int,
global_batch_size: int,
dataloader_batch_size: int,
train_steps: int,
num_workers: int = 8,
drop_last: bool = True,
pin_memory: bool = True,
prefetch_factor: int = 2,
seed: int = 0,
build_collate_fn: bool = True,
collate_fn_kwargs: Optional[Dict[str, Any]] = None,
) -> "DistributedDataloader":
if collate_fn_kwargs is None:
collate_fn_kwargs = {}
# TODO: also need dyn_bsz here?
parallel_state = get_parallel_state()
num_micro_batch = global_batch_size // (
micro_batch_size * parallel_state.dp_size
) # num_micro_batch = num accumulation steps
logger.info_rank0(
f"train_steps: {train_steps},"
f"num_micro_batch: {num_micro_batch}, "
f"micro_batch_size: {micro_batch_size}, global_batch_size: {global_batch_size}, "
f"dp_size: {parallel_state.dp_size}, sp_size: {parallel_state.sp_size}."
)
if build_collate_fn:
collate_fn = MainCollator(**collate_fn_kwargs)
else:
collate_fn = NoopDataCollator()
collate_fn = MakeMicroBatchCollator(num_micro_batch=num_micro_batch, internal_data_collator=collate_fn)
sampler = None
if not isinstance(dataset, IterableDataset):
sampler = StatefulDistributedSampler(
dataset,
num_replicas=parallel_state.dp_size,
rank=parallel_state.dp_rank,
shuffle=True,
seed=seed,
)
dataloader = DistributedDataloader(
dataset,
batch_size=dataloader_batch_size,
sampler=sampler,
num_workers=num_workers,
collate_fn=collate_fn,
pin_memory=pin_memory,
pin_memory_device=get_device_type(),
drop_last=drop_last,
prefetch_factor=prefetch_factor,
)
return dataloader
|