diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py index 9c552ee..b387917 100644 --- a/ScaFFold/cli.py +++ b/ScaFFold/cli.py @@ -140,6 +140,16 @@ def main(): benchmark_parser.add_argument( "--batch-size", type=int, nargs="+", help="Batch sizes for each volume size." ) + benchmark_parser.add_argument( + "--warmup-batches", + type=int, + help="Number of warmup batches to run per rank before training.", + ) + benchmark_parser.add_argument( + "--dataloader-num-workers", + type=int, + help="Number of DataLoader worker processes per rank.", + ) benchmark_parser.add_argument( "--optimizer", type=str, diff --git a/ScaFFold/configs/benchmark_default.yml b/ScaFFold/configs/benchmark_default.yml index e96b103..2d1d414 100644 --- a/ScaFFold/configs/benchmark_default.yml +++ b/ScaFFold/configs/benchmark_default.yml @@ -8,6 +8,7 @@ problem_scale: 8 # Determines dataset resolution and number of unet_bottleneck_dim: 3 # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8. seed: 42 # Random seed. batch_size: 1 # Batch sizes for each vol size. +dataloader_num_workers: 4 # Number of DataLoader worker processes per rank. optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. dc_num_shards: [1, 1, 2] # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum dc_shard_dims: [2, 3, 4] # DistConv param: dimension on which to shard @@ -29,6 +30,6 @@ framework: "torch" # The DL framework to train with. Only valid checkpoint_dir: "checkpoints" # Subfolder in which to save training checkpoints. loss_freq: 1 # Number of epochs between logging the overall loss. normalize: 1 # Cateogry search normalization parameter -warmup_epochs: 1 # How many warmup epochs before training +warmup_batches: 5 # How many warmup batches per rank to run before training. dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse. -target_dice: 0.95 \ No newline at end of file +target_dice: 0.95 diff --git a/ScaFFold/configs/benchmark_testing.yml b/ScaFFold/configs/benchmark_testing.yml index 5fea4d0..f37749c 100644 --- a/ScaFFold/configs/benchmark_testing.yml +++ b/ScaFFold/configs/benchmark_testing.yml @@ -8,6 +8,7 @@ problem_scale: 6 # Determines dataset resolution and number of unet_bottleneck_dim: 3 # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8. seed: 42 # Random seed. batch_size: 1 # Batch sizes for each vol size. +dataloader_num_workers: 4 # Number of DataLoader worker processes per rank. optimizer: "ADAM" # "ADAM" is preferred option, otherwise training defautls to RMSProp. num_shards: [1, 1, 1] # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum shard_dim: [2, 3, 4] # DistConv param: dimension on which to shard @@ -29,6 +30,6 @@ framework: "torch" # The DL framework to train with. Only valid checkpoint_dir: "checkpoints" # Subfolder in which to save training checkpoints. loss_freq: 1 # Number of epochs between logging the overall loss. normalize: 1 # Cateogry search normalization parameter -warmup_epochs: 1 # How many warmup epochs before training +warmup_batches: 5 # How many warmup batches per rank to run before training. dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse. target_dice: 0.95 diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py index fc19f8c..2e74abe 100644 --- a/ScaFFold/datagen/get_dataset.py +++ b/ScaFFold/datagen/get_dataset.py @@ -29,7 +29,9 @@ from ScaFFold.datagen import volumegen META_FILENAME = "meta.yaml" +DATASET_FORMAT_VERSION = 2 INCLUDE_KEYS = [ + "dataset_format_version", "n_categories", "n_instances_used_per_fractal", "problem_scale", @@ -116,8 +118,10 @@ def get_dataset( root.mkdir(exist_ok=True) # Get dict of required keys and compute config_id + config_dict = vars(config).copy() + config_dict["dataset_format_version"] = DATASET_FORMAT_VERSION volume_config = _get_required_keys_dict( - config=vars(config), include_keys=INCLUDE_KEYS + config=config_dict, include_keys=INCLUDE_KEYS ) config_id = _hash_volume_config(volume_config) commit = _git_commit_short() @@ -136,6 +140,8 @@ def get_dataset( meta = yaml.safe_load(meta_path.read_text()) if meta.get("config_id") != config_id: continue + if meta.get("dataset_format_version", 1) != DATASET_FORMAT_VERSION: + continue if require_commit and meta.get("code_commit") != commit: continue # If we pass the above checks, this dataset can be reused @@ -186,6 +192,7 @@ def get_dataset( # Write to tmp, then move, so readers never see half-written dataset meta = { "config_id": config_id, + "dataset_format_version": DATASET_FORMAT_VERSION, "config_subset": volume_config, "include_keys": INCLUDE_KEYS, "code_commit": commit, diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py index 479e67e..59dae15 100644 --- a/ScaFFold/datagen/volumegen.py +++ b/ScaFFold/datagen/volumegen.py @@ -180,7 +180,7 @@ def main(config: Dict): dtype=np.float32, ) mask = np.full( - (config.vol_size, config.vol_size, config.vol_size), 0, dtype=np.short + (config.vol_size, config.vol_size, config.vol_size), 0, dtype=np.int64 ) global_vol_idx = curr_vol[0] @@ -223,14 +223,18 @@ def main(config: Dict): # Determine destination folder subdir = "validation" if global_vol_idx in val_indices else "training" + volume_to_save = np.ascontiguousarray( + volume.transpose((3, 0, 1, 2)), dtype=np.float32 + ) + mask_to_save = np.ascontiguousarray(mask, dtype=np.int64) vol_file = os.path.join(vol_path, subdir, f"{global_vol_idx}.npy") with open(vol_file, "wb") as f: - np.save(f, volume) + np.save(f, volume_to_save) mask_file = os.path.join(mask_path, subdir, f"{global_vol_idx}_mask.npy") with open(mask_file, "wb") as f: - np.save(f, mask) + np.save(f, mask_to_save) end_time = time.time() total_time = end_time - start_time diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py index 50a198d..7e3c5ff 100644 --- a/ScaFFold/utils/config_utils.py +++ b/ScaFFold/utils/config_utils.py @@ -50,6 +50,7 @@ def __init__(self, config_dict): self.n_instances_used_per_fractal = config_dict["n_instances_used_per_fractal"] self.scale = 1 self.batch_size = config_dict["batch_size"] + self.dataloader_num_workers = config_dict.get("dataloader_num_workers", 4) self.epochs = config_dict["epochs"] self.optimizer = config_dict["optimizer"] self.disable_scheduler = bool(config_dict["disable_scheduler"]) @@ -66,7 +67,7 @@ def __init__(self, config_dict): self.loss_freq = config_dict["loss_freq"] self.checkpoint_dir = config_dict["checkpoint_dir"] self.normalize = config_dict["normalize"] - self.warmup_epochs = config_dict["warmup_epochs"] + self.warmup_batches = config_dict.get("warmup_batches") self.dataset_reuse_enforce_commit_id = config_dict[ "dataset_reuse_enforce_commit_id" ] diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py index 725854c..8578093 100644 --- a/ScaFFold/utils/data_loading.py +++ b/ScaFFold/utils/data_loading.py @@ -19,10 +19,15 @@ import numpy as np import torch +import yaml from torch.utils.data import Dataset from ScaFFold.utils.utils import customlog +DATASET_FORMAT_VERSION = 2 +LEGACY_DATASET_FORMAT_VERSION = 1 +META_FILENAME = "meta.yaml" + class BasicDataset(Dataset): def __init__( @@ -31,6 +36,8 @@ def __init__( self.images_dir = Path(images_dir) self.mask_dir = Path(mask_dir) self.mask_suffix = mask_suffix + self.dataset_root = self.images_dir.parents[1] + self.dataset_format_version = self._load_dataset_format_version() self.ids = [ splitext(file)[0] @@ -49,25 +56,56 @@ def __init__( data = pickle.load(data_file) self.mask_values = data["mask_values"] customlog(f"Unique mask values: {self.mask_values}") + customlog(f"Dataset format version: {self.dataset_format_version}") def __len__(self): return len(self.ids) @staticmethod - def preprocess(mask_values, img, is_mask): - if is_mask: - mask = np.zeros((img.shape[0], img.shape[1], img.shape[2]), dtype=np.short) - for i, v in enumerate(mask_values): - if img.ndim == 3: - mask[img == v] = i - else: - mask[(img == v).all(-1)] = i + def _load_numpy_array(path): + with open(path, "rb") as handle: + return np.load(handle) + + def _load_dataset_format_version(self): + meta_path = self.dataset_root / META_FILENAME + if not meta_path.exists(): + return LEGACY_DATASET_FORMAT_VERSION + + try: + with open(meta_path, "r") as meta_file: + meta = yaml.safe_load(meta_file) or {} + except Exception as exc: + customlog( + f"Failed to read dataset metadata from {meta_path}: {exc}. Falling back to legacy loader." + ) + return LEGACY_DATASET_FORMAT_VERSION - return mask + return int(meta.get("dataset_format_version", LEGACY_DATASET_FORMAT_VERSION)) - else: - img = img.transpose((3, 0, 1, 2)) - return img + @staticmethod + def _prepare_legacy_image(img): + return np.ascontiguousarray(img.transpose((3, 0, 1, 2)), dtype=np.float32) + + @staticmethod + def _prepare_legacy_mask(mask_values, mask): + remapped = np.zeros( + (mask.shape[0], mask.shape[1], mask.shape[2]), dtype=np.int64 + ) + for i, value in enumerate(mask_values): + if mask.ndim == 3: + remapped[mask == value] = i + else: + remapped[(mask == value).all(-1)] = i + + return remapped + + @staticmethod + def _prepare_optimized_image(img): + return np.ascontiguousarray(img, dtype=np.float32) + + @staticmethod + def _prepare_optimized_mask(mask): + return np.ascontiguousarray(mask, dtype=np.int64) def __getitem__(self, idx): name = self.ids[idx] @@ -80,19 +118,19 @@ def __getitem__(self, idx): assert len(mask_file) == 1, ( f"Either no mask or multiple masks found for the ID {name}: {mask_file}" ) - with open(mask_file[0], "rb") as f: - mask = np.load(f) - f.close() - with open(img_file[0], "rb") as f: - img = np.load(f) - f.close() + mask = self._load_numpy_array(mask_file[0]) + img = self._load_numpy_array(img_file[0]) - img = self.preprocess(self.mask_values, img, is_mask=False) - mask = self.preprocess(self.mask_values, mask, is_mask=True) + if self.dataset_format_version >= DATASET_FORMAT_VERSION: + img = self._prepare_optimized_image(img) + mask = self._prepare_optimized_mask(mask) + else: + img = self._prepare_legacy_image(img) + mask = self._prepare_legacy_mask(self.mask_values, mask) return { - "image": torch.as_tensor(img.copy()).float().contiguous(), - "mask": torch.as_tensor(mask.copy()).long().contiguous(), + "image": torch.from_numpy(img).contiguous().float(), + "mask": torch.from_numpy(mask).contiguous().long(), } diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index 20f9251..9151e9d 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -29,7 +29,7 @@ import torch.nn.functional as F from distconv import DCTensor from torch import optim -from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor +from torch.distributed.tensor import DTensor from torch.utils.data import DataLoader from tqdm import tqdm @@ -73,6 +73,9 @@ def __init__(self, model, config, device, log): self.criterion = None self.global_step = 0 self.start_epoch = -1 + self.ps = None # DistConv ParallelStrategy + self.spatial_mesh = None # Spatial mesh for use w/ DistConv + self.ddp_placements = None # DDP placements for use w/ DistConv self.checkpoint_path_absolute = str( self.config.run_dir + "/" + self.config.checkpoint_dir @@ -130,11 +133,17 @@ def create_dataloaders(self): self.create_dataset() self.create_sampler() + num_workers = self.config.dataloader_num_workers loader_args = dict( - batch_size=self.config.batch_size, num_workers=1, pin_memory=True + batch_size=self.config.batch_size, + num_workers=num_workers, + pin_memory=True, ) + if num_workers > 0: + loader_args["persistent_workers"] = True + loader_args["prefetch_factor"] = 2 self.log.debug( - f"dataloader num_workers={loader_args['num_workers']}, os.cpu_count()={os.cpu_count()}, self.world_size={self.world_size} " + f"dataloader num_workers={loader_args['num_workers']}, prefetch_factor={loader_args.get('prefetch_factor')}, persistent_workers={loader_args.get('persistent_workers', False)}, os.cpu_count()={os.cpu_count()}, self.world_size={self.world_size} " ) self.train_loader = DataLoader( self.train_set, sampler=self.train_sampler, **loader_args @@ -205,6 +214,10 @@ def __init__(self, model, config, device, log): async_save=getattr(self.config, "async_save", False), ) + self.ps = None # DistConv ParallelStrategy + self.spatial_mesh = None # Spatial mesh for use w/ DistConv + self.ddp_placements = None # DDP placements for use w/ DistConv + def cleanup_or_resume(self): """ Clean up existing train stats and checkpoints, @@ -326,178 +339,159 @@ def _get_memsize(self, tensor, tensor_label: str, verbosity: int = 0): tensor_memory_gb = tensor_memory_bytes / (1024**3) self.log.info(f"{tensor_label} size on GPU: {tensor_memory_gb:.2f} GB") - def train(self): - """ - Execute model training - """ - - self.cleanup_or_resume() + def warmup(self): + """Run warmup iterations before the main training loop.""" + warmup_batches = self.config.warmup_batches + if warmup_batches <= 0: + return - # DistConv ParallelStrategy - ps = getattr(self.config, "_parallel_strategy", None) - if ps is None: - raise RuntimeError( - "ParallelStrategy not found in config. Set config._parallel_strategy when wrapping model with DistConvDDP." + if self.config.dist: + self.train_loader.sampler.set_epoch(0) + + # Match the main training path as closely as possible. + self.model.train() + self.optimizer.zero_grad(set_to_none=False) + start_warmup = time.time() + max_batches = min(warmup_batches, len(self.train_loader)) + self.log.info(f"Running {max_batches} warmup batch(es) per rank") + + for batch_idx, batch in enumerate(self.train_loader): + if batch_idx >= max_batches: + break + + images, true_masks = batch["image"], batch["mask"] + + images = images.to( + device=self.device, + dtype=torch.float32, + memory_format=torch.channels_last_3d, + non_blocking=True, ) - # Get the process group for spatial sharding mesh - spatial_mesh = ps.device_mesh[ps.distconv_dim_names] - - # Get placements for DDP sharding - num_spatial_dims = len(ps.shard_dim) - ddp_placements = [Shard(0)] + [Replicate()] * num_spatial_dims - - warmup_epochs = self.config.warmup_epochs - if warmup_epochs > 0: - begin_code_region("warmup") - # Keep BN/Dropout from changing behavior/statistics - self.model.train() - start_warmup = time.time() - self.log.info(f"Running {warmup_epochs} warmup epoch(s)") - - for _ in range(warmup_epochs): - for i, batch in enumerate(self.train_loader): - self.log.debug(f" warmup: batch {i} / {len(self.train_loader)}") - batch_t_start = time.time() - # Load initial samples and labels - images, true_masks = batch["image"], batch["mask"] - - # Move samples and labels to GPU - images = images.to( - device=self.device, - dtype=torch.float32, - memory_format=torch.channels_last_3d, - non_blocking=True, - ) - self._get_memsize(images, "Original image", self.config.verbose) - true_masks = true_masks.to( - device=self.device, dtype=torch.long, non_blocking=True - ) - self._get_memsize(images, "Original label", self.config.verbose) - - # Add a dummy channel dimension to get 5D [B, 1, D, H, W] - true_masks = true_masks.unsqueeze(1) - - # Data parallel sharding - images_dp = DTensor.from_local( - images, ps.device_mesh, placements=ddp_placements - ).to_local() - - true_masks_dp = DTensor.from_local( - true_masks, ps.device_mesh, placements=ddp_placements - ).to_local() - - # Delete source tensors immediately after use to keep memory down - del images, true_masks - - # Spatial sharding via DistConv - images_dc = DCTensor.distribute(images_dp, ps) - true_masks_dc = DCTensor.distribute(true_masks_dp, ps) - self._get_memsize(images_dc, "Sharded image", self.config.verbose) - - with torch.autocast( - self.device.type if self.device.type != "mps" else "cpu", - enabled=self.config.torch_amp, - ): - # Forward on DCTensor - self.log.debug(f" warmup: running forward pass") - masks_pred_dc = self.model(images_dc) - self.log.debug(f" warmup: forward pass complete") - - # Extract the underlying PyTorch local tensors - local_preds = masks_pred_dc - local_labels_5d = true_masks_dc - - # Remove the dummy channel dimension so CE Loss is happy [B, D, H, W] - local_labels = local_labels_5d.squeeze(1) - if self.world_rank == 0: - self.log.debug( - f" warmup: Local Preds Shape: {local_preds.shape}" - ) - # Should be something like [1, 6, 128, 128, 64] if sharding Width by 2 - self.log.debug( - f" warmup: Local Labels Shape: {local_labels.shape}" - ) - # Should be something like [1, 128, 128, 64] + true_masks = true_masks.to( + device=self.device, dtype=torch.long, non_blocking=True + ).contiguous() + + # Add a dummy channel dimension to get 5D [B, 1, D, H, W] + true_masks = true_masks.unsqueeze(1) + + # Data parallel sharding + images_dp = DTensor.from_local( + images, self.ps.device_mesh, placements=self.ddp_placements + ).to_local() + + true_masks_dp = DTensor.from_local( + true_masks, self.ps.device_mesh, placements=self.ddp_placements + ).to_local() + + # Spatial sharding via DistConv + images_dc = DCTensor.distribute(images_dp, self.ps) + true_masks_dc = DCTensor.distribute(true_masks_dp, self.ps) + self._get_memsize(images_dc, "Sharded image", self.config.verbose) + + with torch.autocast( + self.device.type if self.device.type != "mps" else "cpu", + enabled=self.config.torch_amp, + ): + # Forward on DCTensor + self.log.debug(f" warmup: running forward pass") + masks_pred_dc = self.model(images_dc) + self.log.debug(f" warmup: forward pass complete") - # --- SHARDED LOSS CALCULATION --- - current_mem = torch.cuda.memory_allocated() / (1024**3) - self.log.debug( - f" warmup: Calculating sharded loss. Mem: {current_mem:.2f} GB." - ) + # Extract the underlying PyTorch local tensors + local_preds = masks_pred_dc + local_labels_5d = true_masks_dc - # 1. Sharded Cross Entropy - local_ce_sum = F.cross_entropy( - local_preds, local_labels, reduction="sum" - ) + # Remove the dummy channel dimension so CE Loss is happy [B, D, H, W] + local_labels = local_labels_5d.squeeze(1) + if self.world_rank == 0: + self.log.debug(f" warmup: Local Preds Shape: {local_preds.shape}") + # Should be something like [1, 6, 128, 128, 64] if sharding Width by 2 + self.log.debug( + f" warmup: Local Labels Shape: {local_labels.shape}" + ) + # Should be something like [1, 128, 128, 64] - # Pass the spatial_mesh directly - global_ce_sum = SpatialAllReduce.apply( - local_ce_sum, spatial_mesh - ) + # --- SHARDED LOSS CALCULATION --- + current_mem = torch.cuda.memory_allocated() / (1024**3) + self.log.debug( + f" warmup: Calculating sharded loss. Mem: {current_mem:.2f} GB." + ) - global_total_voxels = local_labels.numel() * math.prod( - self.config.dc_num_shards - ) - loss_ce = global_ce_sum / global_total_voxels + # 1. Sharded Cross Entropy + local_ce_sum = F.cross_entropy( + local_preds, local_labels, reduction="sum" + ) - # 2. Sharded Dice Loss - local_preds_softmax = F.softmax(local_preds, dim=1).float() - local_labels_one_hot = ( - F.one_hot( - local_labels, num_classes=self.config.n_categories + 1 - ) - .permute(0, 4, 1, 2, 3) - .float() - ) - dice_scores = compute_sharded_dice( - local_preds_softmax, local_labels_one_hot, spatial_mesh - ) - loss_dice = 1.0 - dice_scores.mean() + # Pass the spatial_mesh directly + global_ce_sum = SpatialAllReduce.apply(local_ce_sum, self.spatial_mesh) - # 3. Combine Loss - loss = loss_ce + loss_dice + global_total_voxels = local_labels.numel() * math.prod( + self.config.dc_num_shards + ) + loss_ce = global_ce_sum / global_total_voxels + + # 2. Sharded Dice Loss + local_preds_softmax = F.softmax(local_preds, dim=1).float() + local_labels_one_hot = ( + F.one_hot(local_labels, num_classes=self.config.n_categories + 1) + .permute(0, 4, 1, 2, 3) + .float() + ) + dice_scores = compute_sharded_dice( + local_preds_softmax, local_labels_one_hot, self.spatial_mesh + ) + loss_dice = 1.0 - dice_scores.mean() - self.log.debug( - f" warmup: loss calculation complete. Proceeding to backward pass" - ) + # 3. Combine Loss + loss = loss_ce + loss_dice - # Backward pass - self.grad_scaler.scale(loss).backward() - self.log.debug( - f" warmup: backward pass complete. Stepping optimizer" - ) + self.log.debug( + f" warmup: loss calculation complete. Proceeding to backward pass" + ) - self.grad_scaler.step(self.optimizer) - self.grad_scaler.update() + # Backward pass + self.grad_scaler.scale(loss).backward() + self.grad_scaler.unscale_(self.optimizer) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) + self.log.debug(f" warmup: backward pass complete. Stepping optimizer") + + self.grad_scaler.step(self.optimizer) + self.grad_scaler.update() + + # Free memory aggressively + del images_dc, true_masks_dc, masks_pred_dc + del ( + local_preds, + local_labels, + local_preds_softmax, + local_labels_one_hot, + ) + del loss_ce, loss_dice, loss, images_dp, true_masks_dp - # Free memory aggressively - del images_dc, true_masks_dc, masks_pred_dc - del ( - local_preds, - local_labels, - local_preds_softmax, - local_labels_one_hot, - ) - del loss_ce, loss_dice, loss, images_dp, true_masks_dp + if self.world_rank == 0: + peak_alloc = torch.cuda.max_memory_allocated() / (1024**3) + peak_reserved = torch.cuda.max_memory_reserved() / (1024**3) + self.log.debug( + f"[MEM-PEAK] Peak alloc: {peak_alloc:.2f} GiB | Peak reserved: {peak_reserved:.2f} GiB", + ) + batch_t_end = time.time() + self.log.debug( + f" warmup: batch {batch_idx} completed in {batch_t_end - start_warmup} seconds" + ) - if self.world_rank == 0: - peak_alloc = torch.cuda.max_memory_allocated() / (1024**3) - peak_reserved = torch.cuda.max_memory_reserved() / (1024**3) - self.log.debug( - f"[MEM-PEAK] Peak alloc: {peak_alloc:.2f} GiB | Peak reserved: {peak_reserved:.2f} GiB", - ) - batch_t_end = time.time() - self.log.debug( - f" warmup: batch {i} completed in {batch_t_end - batch_t_start} seconds" - ) + # Nuke any accumulated grads so the first real step starts clean + for p in self.model.parameters(): + p.grad = None + self.optimizer.zero_grad(set_to_none=True) - # Nuke any accumulated grads so the first real step starts clean - for p in self.model.parameters(): - p.grad = None - self.optimizer.zero_grad(set_to_none=True) + if self.config.dist: torch.distributed.barrier() - end_code_region("warmup") - self.log.info(f"Done warmup. Took {int(time.time() - start_warmup)}s") + self.log.info(f"Done warmup. Took {int(time.time() - start_warmup)}s") + + def train(self): + """ + Execute model training + """ epoch = 1 dice_score_train = 0 @@ -512,9 +506,7 @@ def train(self): # Timer and tracking variables epoch_start_time = time.time() - train_dice_curr = 0 train_dice_total = 0 - CE_loss = 0 epoch_loss = 0 # Accumulator for per-batch losses # Set necessary modes/states @@ -522,6 +514,7 @@ def train(self): self.train_loader.sampler.set_epoch(epoch) self.val_loader.sampler.set_epoch(epoch) self.model.train() + self.optimizer.zero_grad(set_to_none=False) estr = ( f"{epoch}" @@ -535,8 +528,6 @@ def train(self): unit="img", disable=True if self.world_rank != 0 else False, ) as pbar: - batch_step = 0 - begin_code_region("batch_loop") for batch in self.train_loader: # Load initial samples and labels @@ -560,19 +551,21 @@ def train(self): # Data parallel sharding images_dp = DTensor.from_local( - images, ps.device_mesh, placements=ddp_placements + images, self.ps.device_mesh, placements=self.ddp_placements ).to_local() true_masks_dp = DTensor.from_local( - true_masks, ps.device_mesh, placements=ddp_placements + true_masks, + self.ps.device_mesh, + placements=self.ddp_placements, ).to_local() # Delete source tensors immediately after use to keep memory down del images, true_masks # Spatial sharding via DistConv - images_dc = DCTensor.distribute(images_dp, ps) - true_masks_dc = DCTensor.distribute(true_masks_dp, ps) + images_dc = DCTensor.distribute(images_dp, self.ps) + true_masks_dc = DCTensor.distribute(true_masks_dp, self.ps) self._get_memsize( images_dc, "Sharded image", self.config.verbose ) @@ -619,7 +612,7 @@ def train(self): # Pass the spatial_mesh directly global_ce_sum = SpatialAllReduce.apply( - local_ce_sum, spatial_mesh + local_ce_sum, self.spatial_mesh ) global_total_voxels = local_labels.numel() * math.prod( @@ -640,7 +633,9 @@ def train(self): # Compute sharded dice using new function dice_scores = compute_sharded_dice( - local_preds_softmax, local_labels_one_hot, spatial_mesh + local_preds_softmax, + local_labels_one_hot, + self.spatial_mesh, ) loss_dice = 1.0 - dice_scores.mean() @@ -657,23 +652,20 @@ def train(self): gather_and_print_mem(self.log, "post_backward") begin_code_region("step_and_update") - if batch_step + 1 == len(self.train_loader): - self.grad_scaler.unscale_(self.optimizer) - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), max_norm=1.0 - ) - self.grad_scaler.step(self.optimizer) - gather_and_print_mem(self.log, "after_optim_step") - - self.grad_scaler.update() - self.optimizer.zero_grad(set_to_none=False) + self.grad_scaler.unscale_(self.optimizer) + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), max_norm=1.0 + ) + self.grad_scaler.step(self.optimizer) + gather_and_print_mem(self.log, "after_optim_step") + self.grad_scaler.update() + self.optimizer.zero_grad(set_to_none=False) end_code_region("step_and_update") # Update the loss begin_code_region("update_loss") pbar.update(images_dc.shape[0]) self.global_step += 1 - batch_step += 1 # Stay on GPU epoch_loss += loss.detach() end_code_region("update_loss") diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py index a11a8c3..ab20c4e 100644 --- a/ScaFFold/worker.py +++ b/ScaFFold/worker.py @@ -24,8 +24,8 @@ import torch import torch.distributed as dist import yaml -from distconv import DCTensor, DistConvDDP, ParallelStrategy -from torch.nn.parallel import DistributedDataParallel as DDP +from distconv import DistConvDDP, ParallelStrategy +from torch.distributed.tensor import Replicate, Shard from ScaFFold.datagen.get_dataset import get_dataset from ScaFFold.unet import UNet @@ -214,6 +214,11 @@ def main(kwargs_dict: dict = {}): torch.backends.cudnn.benchmark = False torch.use_deterministic_algorithms(True, warn_only=True) trainer = PyTorchTrainer(model, config, device, log) + trainer.ps = ps + trainer.spatial_mesh = ps.device_mesh[ps.distconv_dim_names] + num_spatial_dims = len(ps.shard_dim) + trainer.ddp_placements = [Shard(0)] + [Replicate()] * num_spatial_dims + else: raise RuntimeError( "Invalid framework specified. Currently [torch] is the supported framework." @@ -225,6 +230,12 @@ def main(kwargs_dict: dict = {}): ranks_per_node = get_local_size() prof_ctx, TORCH_PERF_LOCAL = get_torch_context(ranks_per_node, rank) with prof_ctx as prof: + begin_code_region("cleanup_or_resume") + trainer.cleanup_or_resume() + end_code_region("cleanup_or_resume") + begin_code_region("warmup") + trainer.warmup() + end_code_region("warmup") begin_code_region("train") trainer.train() end_code_region("train")