Merge branch 'vision-merge' into 'main' (539e39b6) · Commits · candle / Megatron-LM

megatron/arguments.py

+14 −2

Original line number	Diff line number	Diff line
		@@ -518,6 +518,9 @@ def _add_initialization_args(parser):
		group.add_argument('--seed', type=int, default=1234,
		help='Random seed used for python, numpy, '
		'pytorch, and cuda.')
		group.add_argument('--data-parallel-random-init', action='store_true',
		help='Enable random initialization of params '
		'across data parallel ranks')
		group.add_argument('--init-method-std', type=float, default=0.02,
		help='Standard deviation of the zero mean normal '
		'distribution used for weight initialization.')
		@@ -835,11 +838,20 @@ def _add_vit_args(parser):

		group.add_argument('--num-classes', type=int, default=1000,
		help='num of classes in vision classificaiton task')
		group.add_argument('--img-dim', type=int, default=224,
		help='Image size for vision classification task')
		group.add_argument('--img-h', type=int, default=224,
		help='Image height for vision classification task')
		group.add_argument('--img-w', type=int, default=224,
		help='Image height for vision classification task')
		group.add_argument('--num-channels', type=int, default=3,
		help='Number of channels in input image data')
		group.add_argument('--patch-dim', type=int, default=16,
		help='patch dimension used in vit')
		group.add_argument('--classes-fraction', type=float, default=1.0,
		help='training with fraction of classes.')
		group.add_argument('--data-per-class-fraction', type=float, default=1.0,
		help='training with fraction of data per class.')
		group.add_argument('--no-data-sharding', action='store_false',
		help='Disable data sharding.',
		dest='data_sharding')

		return parser

megatron/checkpointing.py

+58 −16

Original line number	Diff line number	Diff line
		@@ -65,6 +65,8 @@ def check_checkpoint_args(checkpoint_args):
		_compare('make_vocab_size_divisible_by')
		_compare('padded_vocab_size')
		_compare('tokenizer_type')
		if args.data_parallel_random_init:
		_compare('data_parallel_random_init')
		if get_checkpoint_version() < 3.0:
		_compare('tensor_model_parallel_size',
		old_arg_name='model_parallel_size')
		@@ -72,7 +74,6 @@ def check_checkpoint_args(checkpoint_args):
		_compare('tensor_model_parallel_size')
		_compare('pipeline_model_parallel_size')


		def ensure_directory_exists(filename):
		"""Build filename's path if it does not already exists."""
		dirname = os.path.dirname(filename)
		@@ -140,6 +141,32 @@ def read_metadata(tracker_filename):
		return max_iter, release


		def get_rng_state():
		""" collect rng state across data parallel ranks """
		args = get_args()
		rng_state = {
		'random_rng_state': random.getstate(),
		'np_rng_state': np.random.get_state(),
		'torch_rng_state': torch.get_rng_state(),
		'cuda_rng_state': torch.cuda.get_rng_state(),
		'rng_tracker_states': mpu.get_cuda_rng_tracker().get_states()}

		rng_state_list = None
		if torch.distributed.is_initialized() and \
		mpu.get_data_parallel_world_size() > 1 and \
		args.data_parallel_random_init:
		rng_state_list = \
		[None for i in range(mpu.get_data_parallel_world_size())]
		torch.distributed.all_gather_object(
		rng_state_list,
		rng_state,
		group=mpu.get_data_parallel_group())
		else:
		rng_state_list = [rng_state]

		return rng_state_list


		def save_checkpoint(iteration, model, optimizer, lr_scheduler):
		"""Save a model checkpoint."""
		args = get_args()
		@@ -150,6 +177,9 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):
		print_rank_0('saving checkpoint at iteration {:7d} to {}'.format(
		iteration, args.save))

		# collect rng state across data parallel ranks
		rng_state = get_rng_state()

		if not torch.distributed.is_initialized() or mpu.get_data_parallel_rank() == 0:

		# Arguments, iteration, and model.
		@@ -173,12 +203,7 @@ def save_checkpoint(iteration, model, optimizer, lr_scheduler):

		# RNG states.
		if not args.no_save_rng:
		state_dict['random_rng_state'] = random.getstate()
		state_dict['np_rng_state'] = np.random.get_state()
		state_dict['torch_rng_state'] = torch.get_rng_state()
		state_dict['cuda_rng_state'] = torch.cuda.get_rng_state()
		state_dict['rng_tracker_states'] \
		= mpu.get_cuda_rng_tracker().get_states()
		state_dict["rng_state"] = rng_state

		# Save.
		checkpoint_name = get_checkpoint_name(args.save, iteration)
		@@ -381,6 +406,23 @@ def load_checkpoint(model, optimizer, lr_scheduler, load_arg='load', strict=True
		# rng states.
		if not release and not args.finetune and not args.no_load_rng:
		try:
		if 'rng_state' in state_dict:
		# access rng_state for data parallel rank
		if args.data_parallel_random_init:

		rng_state = state_dict['rng_state'][mpu.get_data_parallel_rank()]
		else:
		rng_state = state_dict['rng_state'][0]
		random.setstate(rng_state['random_rng_state'])
		np.random.set_state(rng_state['np_rng_state'])
		torch.set_rng_state(rng_state['torch_rng_state'])
		torch.cuda.set_rng_state(rng_state['cuda_rng_state'])
		# Check for empty states array
		if not rng_state['rng_tracker_states']:
		raise KeyError
		mpu.get_cuda_rng_tracker().set_states(
		rng_state['rng_tracker_states'])
		else: # backward compatability
		random.setstate(state_dict['random_rng_state'])
		np.random.set_state(state_dict['np_rng_state'])
		torch.set_rng_state(state_dict['torch_rng_state'])

megatron/data/data_samplers.py

+55 −13

Original line number	Diff line number	Diff line
		@@ -16,8 +16,10 @@
		"""Dataloaders."""


		import torch
		import random
		import torch
		import numpy as np
		from torch.utils.data import Dataset
		from megatron import get_args
		from megatron import mpu

		@@ -39,11 +41,13 @@ def build_pretraining_data_loader(dataset, consumed_samples):
		data_parallel_size=mpu.get_data_parallel_world_size())
		elif args.dataloader_type == 'cyclic':
		batch_sampler = MegatronPretrainingRandomSampler(
		dataset,
		total_samples=len(dataset),
		consumed_samples=consumed_samples,
		micro_batch_size=args.micro_batch_size,
		data_parallel_rank=mpu.get_data_parallel_rank(),
		data_parallel_size=mpu.get_data_parallel_world_size())
		data_parallel_size=mpu.get_data_parallel_world_size(),
		data_sharding=args.data_sharding)
		else:
		raise Exception('{} dataloader type is not supported.'.format(
		args.dataloader_type))
		@@ -103,16 +107,40 @@ class MegatronPretrainingSampler:
		yield batch[start_idx:end_idx]


		class RandomSeedDataset(Dataset):

		def __init__(self, dataset):
		args = get_args()
		self.base_seed = args.seed
		self.curr_seed = args.seed
		self.dataset = dataset

		def __len__(self):
		return len(self.dataset)

		def set_epoch(self, epoch):
		self.curr_seed = self.base_seed + epoch

		def __getitem__(self, idx):
		seed = idx + self.curr_seed
		torch.manual_seed(seed)
		random.seed(seed)
		np.random.seed(seed)
		return self.dataset[idx]


		class MegatronPretrainingRandomSampler:

		def __init__(self, total_samples, consumed_samples, micro_batch_size,
		data_parallel_rank, data_parallel_size):
		def __init__(self, dataset, total_samples, consumed_samples, micro_batch_size,
		data_parallel_rank, data_parallel_size, data_sharding):
		# Keep a copy of input params for later use.
		self.dataset = dataset
		self.total_samples = total_samples
		self.consumed_samples = consumed_samples
		self.micro_batch_size = micro_batch_size
		self.data_parallel_rank = data_parallel_rank
		self.data_parallel_size = data_parallel_size
		self.data_sharding = data_sharding
		self.micro_batch_times_data_parallel_size = \
		self.micro_batch_size * data_parallel_size
		self.last_batch_size = \
		@@ -136,7 +164,11 @@ class MegatronPretrainingRandomSampler:
		current_epoch_samples = self.consumed_samples % active_total_samples
		assert current_epoch_samples % self.micro_batch_times_data_parallel_size == 0

		if isinstance(self.dataset, RandomSeedDataset):
		self.dataset.set_epoch(self.epoch)

		# data sharding and random sampling
		if self.data_sharding:
		bucket_size = (self.total_samples // self.micro_batch_times_data_parallel_size) \
		* self.micro_batch_size
		bucket_offset = current_epoch_samples // self.data_parallel_size
		@@ -146,6 +178,16 @@ class MegatronPretrainingRandomSampler:
		g.manual_seed(self.epoch)
		random_idx = torch.randperm(bucket_size, generator=g).tolist()
		idx_range = [start_idx + x for x in random_idx[bucket_offset:]]
		else:
		full_bucket_size = (self.total_samples // self.micro_batch_size) \
		* self.micro_batch_size
		full_bucket_offset = current_epoch_samples
		g = torch.Generator()
		g.manual_seed(self.epoch)
		idx_range_total = \
		torch.randperm(full_bucket_size, generator=g).tolist()
		idx_range_active = idx_range_total[full_bucket_offset:]
		idx_range = idx_range_active[self.data_parallel_rank::self.data_parallel_size]

		batch = []
		# Last batch if not complete will be dropped.

megatron/data/image_folder.py

0 → 100644

+302 −0

Original line number	Diff line number	Diff line
		# BSD 3-Clause License
		#
		# Copyright (c) Soumith Chintala 2016,
		# All rights reserved.
		#
		# Redistribution and use in source and binary forms, with or without
		# modification, are permitted provided that the following conditions are met:
		#
		# * Redistributions of source code must retain the above copyright notice, this
		# list of conditions and the following disclaimer.
		#
		# * Redistributions in binary form must reproduce the above copyright notice,
		# this list of conditions and the following disclaimer in the documentation
		# and/or other materials provided with the distribution.
		#
		# * Neither the name of the copyright holder nor the names of its
		# contributors may be used to endorse or promote products derived from
		# this software without specific prior written permission.

		# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
		# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
		# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
		# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
		# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
		# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
		# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
		# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
		# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
		# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

		# code taken from
		# https://github.com/pytorch/vision/blob/main/torchvision/datasets/folder.py
		# added support for classes_fraction and data_per_class_fraction

		from torchvision.datasets import VisionDataset
		from PIL import Image

		import os
		import os.path
		from typing import Any, Callable, cast, Dict, List, Optional, Tuple
		import numpy as np

		def has_file_allowed_extension(filename: str, extensions: Tuple[str, ...]) -> bool:
		"""Checks if a file is an allowed extension.
		Args:
		filename (string): path to a file
		extensions (tuple of strings): extensions to consider (lowercase)
		Returns:
		bool: True if the filename ends with one of given extensions
		"""
		return filename.lower().endswith(extensions)


		def is_image_file(filename: str) -> bool:
		"""Checks if a file is an allowed image extension.
		Args:
		filename (string): path to a file
		Returns:
		bool: True if the filename ends with a known image extension
		"""
		return has_file_allowed_extension(filename, IMG_EXTENSIONS)


		def make_dataset(
		directory: str,
		class_to_idx: Dict[str, int],
		data_per_class_fraction: float,
		extensions: Optional[Tuple[str, ...]] = None,
		is_valid_file: Optional[Callable[[str], bool]] = None,
		) -> List[Tuple[str, int]]:
		"""Generates a list of samples of a form (path_to_sample, class).
		Args:
		directory (str): root dataset directory
		class_to_idx (Dict[str, int]): dictionary mapping class name to class index
		extensions (optional): A list of allowed extensions.
		Either extensions or is_valid_file should be passed. Defaults to None.
		is_valid_file (optional): A function that takes path of a file
		and checks if the file is a valid file
		(used to check of corrupt files) both extensions and
		is_valid_file should not be passed. Defaults to None.
		Raises:
		ValueError: In case ``extensions`` and ``is_valid_file`` are None or both are not None.
		Returns:
		List[Tuple[str, int]]: samples of a form (path_to_sample, class)
		"""
		instances = []
		directory = os.path.expanduser(directory)
		both_none = extensions is None and is_valid_file is None
		both_something = extensions is not None and is_valid_file is not None
		if both_none or both_something:
		raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")
		if extensions is not None:
		def is_valid_file(x: str) -> bool:
		return has_file_allowed_extension(x, cast(Tuple[str, ...], extensions))
		is_valid_file = cast(Callable[[str], bool], is_valid_file)
		for target_class in sorted(class_to_idx.keys()):
		class_index = class_to_idx[target_class]
		target_dir = os.path.join(directory, target_class)
		if not os.path.isdir(target_dir):
		continue
		local_instances = []
		for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
		for fname in sorted(fnames):
		path = os.path.join(root, fname)
		if is_valid_file(path):
		item = path, class_index
		local_instances.append(item)

		instances.extend(local_instances[0:int(len(local_instances) * data_per_class_fraction)])

		return instances


		class DatasetFolder(VisionDataset):
		"""A generic data loader where the samples are arranged in this way: ::
		root/class_x/xxx.ext
		root/class_x/xxy.ext
		root/class_x/[...]/xxz.ext
		root/class_y/123.ext
		root/class_y/nsdf3.ext
		root/class_y/[...]/asd932_.ext
		Args:
		root (string): Root directory path.
		loader (callable): A function to load a sample given its path.
		extensions (tuple[string]): A list of allowed extensions.
		both extensions and is_valid_file should not be passed.
		transform (callable, optional): A function/transform that takes in
		a sample and returns a transformed version.
		E.g, ``transforms.RandomCrop`` for images.
		target_transform (callable, optional): A function/transform that takes
		in the target and transforms it.
		is_valid_file (callable, optional): A function that takes path of a file
		and check if the file is a valid file (used to check of corrupt files)
		both extensions and is_valid_file should not be passed.
		Attributes:
		classes (list): List of the class names sorted alphabetically.
		class_to_idx (dict): Dict with items (class_name, class_index).
		samples (list): List of (sample path, class_index) tuples
		targets (list): The class_index value for each image in the dataset
		"""

		def __init__(
		self,
		root: str,
		loader: Callable[[str], Any],
		extensions: Optional[Tuple[str, ...]] = None,
		transform: Optional[Callable] = None,
		target_transform: Optional[Callable] = None,
		classes_fraction=1.0,
		data_per_class_fraction=1.0,
		is_valid_file: Optional[Callable[[str], bool]] = None,
		) -> None:
		super(DatasetFolder, self).__init__(root, transform=transform,
		target_transform=target_transform)
		self.classes_fraction = classes_fraction
		self.data_per_class_fraction = data_per_class_fraction
		classes, class_to_idx = self._find_classes(self.root)
		samples = self.make_dataset(self.root,
		class_to_idx,
		self.data_per_class_fraction,
		extensions,
		is_valid_file)
		if len(samples) == 0:
		msg = "Found 0 files in subfolders of: {}\n".format(self.root)
		if extensions is not None:
		msg += "Supported extensions are: {}".format(",".join(extensions))
		raise RuntimeError(msg)

		self.loader = loader
		self.extensions = extensions
		self.total = len(samples)
		self.classes = classes
		self.class_to_idx = class_to_idx
		self.samples = samples
		self.targets = [s[1] for s in samples]

		@staticmethod
		def make_dataset(
		directory: str,
		class_to_idx: Dict[str, int],
		data_per_class_fraction: float,
		extensions: Optional[Tuple[str, ...]] = None,
		is_valid_file: Optional[Callable[[str], bool]] = None,
		) -> List[Tuple[str, int]]:
		return make_dataset(directory,
		class_to_idx,
		data_per_class_fraction,
		extensions=extensions,
		is_valid_file=is_valid_file)

		def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
		"""
		Finds the class folders in a dataset.
		Args:
		dir (string): Root directory path.
		Returns:
		tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.
		Ensures:
		No class is a subdirectory of another.
		"""
		all_classes = [d.name for d in os.scandir(dir) if d.is_dir()]
		classes = all_classes[0:int(len(all_classes) * self.classes_fraction)]
		classes.sort()
		class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
		return classes, class_to_idx

		def __getitem__(self, index: int) -> Tuple[Any, Any]:
		"""
		Args:
		index (int): Index
		Returns:
		tuple: (sample, target) where target is class_index of the target class.
		"""
		curr_index = index
		for x in range(self.total):
		try:
		path, target = self.samples[curr_index]
		sample = self.loader(path)
		break
		except Exception as e:
		curr_index = np.random.randint(0, self.total)

		if self.transform is not None:
		sample = self.transform(sample)
		if self.target_transform is not None:
		target = self.target_transform(target)

		return sample, target

		def __len__(self) -> int:
		return len(self.samples)


		IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')


		def pil_loader(path: str) -> Image.Image:
		# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
		with open(path, 'rb') as f:
		img = Image.open(f)
		return img.convert('RGB')


		# TODO: specify the return type
		def accimage_loader(path: str) -> Any:
		import accimage
		try:
		return accimage.Image(path)
		except IOError:
		# Potentially a decoding problem, fall back to PIL.Image
		return pil_loader(path)


		def default_loader(path: str) -> Any:
		from torchvision import get_image_backend
		if get_image_backend() == 'accimage':
		return accimage_loader(path)
		else:
		return pil_loader(path)


		class ImageFolder(DatasetFolder):
		"""A generic data loader where the images are arranged in this way: ::
		root/dog/xxx.png
		root/dog/xxy.png
		root/dog/[...]/xxz.png
		root/cat/123.png
		root/cat/nsdf3.png
		root/cat/[...]/asd932_.png
		Args:
		root (string): Root directory path.
		transform (callable, optional): A function/transform that takes in an PIL image
		and returns a transformed version. E.g, ``transforms.RandomCrop``
		target_transform (callable, optional): A function/transform that takes in the
		target and transforms it.
		loader (callable, optional): A function to load an image given its path.
		is_valid_file (callable, optional): A function that takes path of an Image file
		and check if the file is a valid file (used to check of corrupt files)
		Attributes:
		classes (list): List of the class names sorted alphabetically.
		class_to_idx (dict): Dict with items (class_name, class_index).
		imgs (list): List of (image path, class_index) tuples
		"""

		def __init__(
		self,
		root: str,
		transform: Optional[Callable] = None,
		target_transform: Optional[Callable] = None,
		classes_fraction=1.0,
		data_per_class_fraction=1.0,
		loader: Callable[[str], Any] = default_loader,
		is_valid_file: Optional[Callable[[str], bool]] = None,
		):
		super(ImageFolder, self).__init__(root, loader, IMG_EXTENSIONS if is_valid_file is None else None,
		transform=transform,
		target_transform=target_transform,
		classes_fraction=classes_fraction,
		data_per_class_fraction=data_per_class_fraction,
		is_valid_file=is_valid_file)
		self.imgs = self.samples

megatron/data/vit_dataset.py

+52 −31

File changed.

Preview size limit exceeded, changes collapsed.