Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / dummy.py: 44%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Deterministic dummy LM dataset. 

16 

17Each sample's tokens are derived from ``base_seed + idx`` so the same global 

18sample index always yields the same content, independent of DP topology or 

19sampler shuffling. 

20""" 

21import logging 

22from typing import Any, Dict 

23 

24import torch 

25from torch.utils.data import Dataset 

26 

27from hyper_parallel.data.registry import DATASET_REGISTRY 

28 

29 

30logger = logging.getLogger(__name__) 

31 

32 

33class DummyDataset(Dataset): 

34 """Deterministic random-token dataset. 

35 

36 Args: 

37 num_samples: Total dataset length. 

38 seq_length: Tokens per sample. 

39 vocab: Token range upper bound (exclusive). 

40 base_seed: Seed offset; each sample uses ``base_seed + idx``. 

41 """ 

42 

43 def __init__(self, num_samples: int, seq_length: int, vocab: int, base_seed: int = 42) -> None: 

44 self.num_samples = int(num_samples) 

45 self.seq_length = int(seq_length) 

46 self.vocab = int(vocab) 

47 self.base_seed = int(base_seed) 

48 

49 def __len__(self) -> int: 

50 return self.num_samples 

51 

52 def __getitem__(self, idx: int) -> Dict[str, Any]: 

53 g = torch.Generator().manual_seed(self.base_seed + idx) 

54 input_ids = torch.randint(0, self.vocab, (self.seq_length,), generator=g) 

55 return {"input_ids": input_ids, "labels": input_ids.clone()} 

56 

57 

58@DATASET_REGISTRY.register("dummy") 

59def build_dummy(*, base: Any, args: Any, **_: Any) -> DummyDataset: 

60 """Build the dummy LM dataset from trainer config. 

61 

62 Picks ``vocab`` from ``model.config.vocab_size`` when available, else 

63 defaults to 32000 (matches the pre-refactor behaviour in BaseTrainer). 

64 """ 

65 seq_len = args.data.max_seq_len 

66 vocab_size = base.model.config.vocab_size 

67 base_seed = args.train.seed 

68 total_samples = base.state.max_steps * args.train.global_batch_size 

69 ds = DummyDataset(total_samples, seq_len, vocab_size, base_seed=base_seed) 

70 logger.info("Dummy dataset created: %d samples, seq_len=%d", total_samples, seq_len) 

71 return ds