Add --num-workers override to avoid DataLoader worker crashes (use 0 on ROCm)

Made-with: Cursor

Author: Alex J Lennon

training/modules/data_pipeline.py

@@ -699,7 +699,8 @@ def create_data_loaders(config: TrainingConfiguration,
                        data_root: Optional[str] = None,
                        pin_memory: Optional[bool] = None,
                        interactive: bool = True,
-                       no_download: bool = False) -> Tuple[DataLoader, DataLoader, DataLoader]:
+                       no_download: bool = False,
+                       num_workers_override: Optional[int] = None) -> Tuple[DataLoader, DataLoader, DataLoader]:
     """
     Create training, validation, and test data loaders
     
@@ -709,12 +710,14 @@ def create_data_loaders(config: TrainingConfiguration,
         pin_memory: Override pin_memory setting (optional)
         interactive: If False, auto-confirm dataset download/prepare (--yes)
         no_download: If True, use only existing prepared data; fail if any missing (--no-download)
+        num_workers_override: If set, override config.hardware.num_workers (e.g. 0 for ROCm)
         
     Returns:
         Tuple of (train_loader, val_loader, test_loader)
     """
     # Use provided pin_memory override or fall back to config
     use_pin_memory = pin_memory if pin_memory is not None else config.hardware.pin_memory
+    num_workers = num_workers_override if num_workers_override is not None else config.hardware.num_workers
 
     # Create datasets
     train_datasets = []
@@ -756,7 +759,7 @@ def create_data_loaders(config: TrainingConfiguration,
         train_dataset,
         batch_size=config.training.batch_size,
         shuffle=True,
-        num_workers=config.hardware.num_workers,
+        num_workers=num_workers,
         pin_memory=use_pin_memory,
         collate_fn=collate_audio_samples,
         drop_last=True  # Ensure consistent batch sizes
@@ -766,7 +769,7 @@ def create_data_loaders(config: TrainingConfiguration,
         val_dataset,
         batch_size=config.training.batch_size,
         shuffle=False,
-        num_workers=config.hardware.num_workers,
+        num_workers=num_workers,
         pin_memory=use_pin_memory,
         collate_fn=collate_audio_samples,
         drop_last=False
@@ -776,7 +779,7 @@ def create_data_loaders(config: TrainingConfiguration,
         test_dataset,
         batch_size=config.training.batch_size,
         shuffle=False,
-        num_workers=config.hardware.num_workers,
+        num_workers=num_workers,
         pin_memory=use_pin_memory,
         collate_fn=collate_audio_samples,
         drop_last=False

training/train.py

@@ -49,7 +49,7 @@ class TCNTrainer:
     clear error handling and progress monitoring.
     """
 
-    def __init__(self, config: TrainingConfiguration, resume_from: Optional[str] = None, data_root: Optional[str] = None, interactive: bool = True, no_download: bool = False):
+    def __init__(self, config: TrainingConfiguration, resume_from: Optional[str] = None, data_root: Optional[str] = None, interactive: bool = True, no_download: bool = False, num_workers_override: Optional[int] = None):
         """
         Initialize trainer with configuration
         
@@ -59,6 +59,7 @@ def __init__(self, config: TrainingConfiguration, resume_from: Optional[str] = N
             data_root: Root directory for data (optional)
             interactive: If False, auto-confirm dataset download/prepare (--yes)
             no_download: If True, use only existing prepared data; fail if any dataset missing (--no-download)
+            num_workers_override: If set, override config.hardware.num_workers for DataLoader (e.g. 0 for ROCm)
         """
         self.config = config
 
@@ -91,7 +92,13 @@ def __init__(self, config: TrainingConfiguration, resume_from: Optional[str] = N
         # Create data loaders
         self.logger.info("Setting up data pipeline...")
         self.train_loader, self.val_loader, self.test_loader = create_data_loaders(
-            config, data_root=data_root, pin_memory=self.pin_memory, interactive=interactive, no_download=no_download)
+            config,
+            data_root=data_root,
+            pin_memory=self.pin_memory,
+            interactive=interactive,
+            no_download=no_download,
+            num_workers_override=num_workers_override,
+        )
 
         # Create model
         self.logger.info("Creating TCN model...")
@@ -632,6 +639,14 @@ def parse_arguments():
         help='Use only existing prepared data; exit with error if any requested dataset is missing or not prepared (no download or MFA)'
     )
 
+    parser.add_argument(
+        '--num-workers',
+        type=int,
+        default=None,
+        metavar='N',
+        help='Override num_workers for DataLoader (e.g. 0 to avoid worker crashes on ROCm)'
+    )
+    
     return parser.parse_args()
 
 
@@ -655,7 +670,14 @@ def main():
 
     try:
         # Create trainer
-        trainer = TCNTrainer(config, resume_from=args.resume, data_root=args.data_root, interactive=not args.yes, no_download=args.no_download)
+        trainer = TCNTrainer(
+            config,
+            resume_from=args.resume,
+            data_root=args.data_root,
+            interactive=not args.yes,
+            no_download=args.no_download,
+            num_workers_override=args.num_workers,
+        )
 
         if args.test_only:
             # Test only mode