adds optimizations for faster training

config/train_gpt_oss_openhermes_fr_memory_optimized.py

@@ -56,7 +56,7 @@ config = GPTOSSEnhancedCustomConfig(
     # MODEL CONFIGURATION - Memory Optimized for GPT-OSS
     # ============================================================================
     model_name="openai/gpt-oss-20b",
-    max_seq_length=4096,                     # Maximize sequence length for A100 VRAM utilization
+    max_seq_length=2048,                     # Shorter context speeds steps without reducing sample count
     use_flash_attention=True,               # Critical for memory efficiency
     use_gradient_checkpointing=True,        # Essential for memory optimization
     
@@ -115,9 +115,10 @@ config = GPTOSSEnhancedCustomConfig(
     },
     
     # Data loading optimized for throughput
-    dataloader_num_workers=4,                # More workers for faster loading
+    dataloader_num_workers=8,                # More workers for faster loading
     dataloader_pin_memory=True,              # Pin memory for faster host->GPU copies
-    dataloader_prefetch_factor=1,            # Lower prefetch to keep VRAM headroom
+    dataloader_prefetch_factor=2,            # Slightly higher prefetch for throughput
+    dataset_num_proc=8,                      # Parallelize HF datasets map/filter
     
     # Memory management optimizations
     max_memory_per_gpu=None,                 # No explicit memory limit; use as much VRAM as available
@@ -197,6 +198,9 @@ config = GPTOSSEnhancedCustomConfig(
         "min_lr": 2e-6,                     # Explicit absolute floor (matches min_lr above)
         "warmup_steps": None,               # Use warmup_ratio instead
     },
+
+    # Packing to increase token utilization per step (supported by TRL)
+    packing=True,
     
     # ============================================================================
     # MONITORING & HUB INTEGRATION

scripts/training/train_gpt_oss.py

@@ -210,6 +210,13 @@ def build_scheduler_kwargs(config):
 def apply_dataset_filtering(dataset, config):
     """Apply filtering based on configuration"""
     
+    # Parallel workers for datasets ops
+    try:
+        import os as _os
+        num_proc = getattr(config, 'dataset_num_proc', None) or (_os.cpu_count() or 1)
+    except Exception:
+        num_proc = 1
+
     # Filter bad entries if specified
     if getattr(config, 'filter_bad_entries', False):
         bad_entry_field = getattr(config, 'bad_entry_field', 'bad_entry')
@@ -220,17 +227,23 @@ def apply_dataset_filtering(dataset, config):
         
         # Filter out bad entries
         if bad_entry_field in dataset.column_names:
-            dataset = dataset.filter(lambda x: not x.get(bad_entry_field, False))
+            def _keep_not_bad_entry(example, _field=bad_entry_field):
+                return not example.get(_field, False)
+            dataset = dataset.filter(_keep_not_bad_entry, num_proc=num_proc)
             print(f"Filtered {original_size - len(dataset)} bad entries")
         
         # Filter out bad prompts
         if bad_prompt_field in dataset.column_names:
-            dataset = dataset.filter(lambda x: not x.get(bad_prompt_field, False))
+            def _keep_not_bad_prompt(example, _field=bad_prompt_field):
+                return not example.get(_field, False)
+            dataset = dataset.filter(_keep_not_bad_prompt, num_proc=num_proc)
             print(f"Filtered bad prompts, remaining: {len(dataset)} examples")
         
         # Filter out bad responses
         if bad_response_field in dataset.column_names:
-            dataset = dataset.filter(lambda x: not x.get(bad_response_field, False))
+            def _keep_not_bad_response(example, _field=bad_response_field):
+                return not example.get(_field, False)
+            dataset = dataset.filter(_keep_not_bad_response, num_proc=num_proc)
             print(f"Filtered bad responses, remaining: {len(dataset)} examples")
     
     # Apply length filtering
@@ -253,7 +266,7 @@ def apply_dataset_filtering(dataset, config):
             return True
         
         original_size = len(dataset)
-        dataset = dataset.filter(length_filter)
+        dataset = dataset.filter(length_filter, num_proc=num_proc)
         print(f"Length filtering: {original_size} -> {len(dataset)} examples")
     
     # Apply sampling if specified
@@ -293,6 +306,13 @@ def format_gpt_oss_harmony_prompt(prompt: str) -> str:
 def process_dataset_format(dataset, config):
     """Process dataset based on format configuration with exact GPT-OSS Harmony compliance"""
     
+    # Parallel workers for datasets ops
+    try:
+        import os as _os
+        num_proc = getattr(config, 'dataset_num_proc', None) or (_os.cpu_count() or 1)
+    except Exception:
+        num_proc = 1
+
     dataset_format = getattr(config, 'dataset_format', 'openhermes_fr')
     input_field = getattr(config, 'input_field', 'prompt')
     target_field = getattr(config, 'target_field', 'accepted_completion')
@@ -325,7 +345,7 @@ def process_dataset_format(dataset, config):
                 return {"prompt": prompt_val, "chosen": chosen_val, "rejected": rejected_val}
 
             keep_cols = [c for c in ['prompt', 'chosen', 'rejected'] if c in dataset.column_names]
-            dataset = dataset.map(id_map, remove_columns=dataset.column_names if keep_cols else dataset.column_names)
+            dataset = dataset.map(id_map, remove_columns=dataset.column_names if keep_cols else dataset.column_names, num_proc=num_proc)
             return dataset
 
         # Custom preference mapping via configured field names
@@ -341,7 +361,7 @@ def process_dataset_format(dataset, config):
                     return {"prompt": prompt_text, "chosen": chosen_text, "rejected": rejected_text}
                 return {"prompt": prompt_val, "chosen": chosen_val, "rejected": rejected_val}
 
-            dataset = dataset.map(to_pref, remove_columns=dataset.column_names)
+            dataset = dataset.map(to_pref, remove_columns=dataset.column_names, num_proc=num_proc)
             return dataset
 
         # If we reach here, we don't have required fields for DPO
@@ -371,7 +391,7 @@ def process_dataset_format(dataset, config):
                     "output": completion
                 }
         
-        dataset = dataset.map(format_openhermes_fr, remove_columns=dataset.column_names)
+        dataset = dataset.map(format_openhermes_fr, remove_columns=dataset.column_names, num_proc=num_proc)
         
     elif dataset_format == "messages":
         # Process messages format (like HuggingFaceH4/Multilingual-Thinking)
@@ -416,7 +436,7 @@ def process_dataset_format(dataset, config):
             
             return {"text": text}
         
-        dataset = dataset.map(format_messages, remove_columns=dataset.column_names)
+        dataset = dataset.map(format_messages, remove_columns=dataset.column_names, num_proc=num_proc)
         
     elif dataset_format == "text":
         # Process plain text format
@@ -427,7 +447,7 @@ def process_dataset_format(dataset, config):
                 text += "</s>"
             return {"text": text}
         
-        dataset = dataset.map(format_text, remove_columns=dataset.column_names)
+        dataset = dataset.map(format_text, remove_columns=dataset.column_names, num_proc=num_proc)
     
     elif dataset_format == "custom":
         # Custom format - user handles this in their config
@@ -652,6 +672,8 @@ def create_sft_config(config, output_dir):
         "bf16": bf16,
         # Some versions support tf32
         "tf32": tf32 if 'tf32' in TrainingArguments.__init__.__code__.co_varnames else None,
+        # Optimizer (optionally use fused AdamW if available through config)
+        "optim": getattr(config, 'optimizer', 'adamw_torch'),
         # Regularization
         "weight_decay": weight_decay,
         "max_grad_norm": max_grad_norm,
@@ -828,6 +850,10 @@ def train_gpt_oss(config_path, experiment_name, output_dir, trackio_url, trainer
         if "max_seq_length" in sft_params:
             sft_kwargs["max_seq_length"] = getattr(config, 'max_seq_length', 2048)
 
+        # Enable sequence packing if supported by TRL (speeds up token utilization)
+        if "packing" in sft_params:
+            sft_kwargs["packing"] = getattr(config, 'packing', False)
+
         # Remove any None values
         sft_kwargs = {k: v for k, v in sft_kwargs.items() if v is not None}
 

src/dataset_utils.py

@@ -122,12 +122,20 @@ class TrackioDatasetManager:
     
     def save_experiments(self, experiments: List[Dict[str, Any]], commit_message: Optional[str] = None) -> bool:
         """
-        Save a list of experiments to the dataset, preserving data integrity.
-        
+        Save a list of experiments to the dataset using a non-destructive union merge.
+
+        - Loads existing experiments (if any) and builds a union by `experiment_id`.
+        - For overlapping IDs, merges JSON fields:
+          - metrics: concatenates lists and de-duplicates by (step, timestamp) for nested entries
+          - parameters: dict-update (new values override)
+          - artifacts: union with de-dup
+          - logs: concatenation with de-dup
+        - Non-JSON scalar fields from incoming experiments take precedence.
+
         Args:
             experiments (List[Dict[str, Any]]): List of experiment dictionaries
             commit_message (Optional[str]): Custom commit message
-            
+
         Returns:
             bool: True if save was successful, False otherwise
         """
@@ -136,24 +144,120 @@ class TrackioDatasetManager:
                 logger.warning("⚠️ No experiments to save")
                 return False
             
-            # Validate all experiments before saving
-            valid_experiments = []
+            # Helpers
+            def _parse_json_field(value, default):
+                try:
+                    if value is None:
+                        return default
+                    if isinstance(value, str):
+                        return json.loads(value) if value else default
+                    return value
+                except Exception:
+                    return default
+            
+            def _metrics_key(entry: Dict[str, Any]):
+                if isinstance(entry, dict):
+                    return (entry.get('step'), entry.get('timestamp'))
+                return (None, json.dumps(entry, sort_keys=True))
+            
+            # Load existing experiments for union merge
+            existing = {}
+            try:
+                for row in self.load_existing_experiments():
+                    exp_id = row.get('experiment_id')
+                    if exp_id:
+                        existing[exp_id] = row
+            except Exception:
+                existing = {}
+            
+            # Validate and merge
+            merged_map: Dict[str, Dict[str, Any]] = {}
+            # Seed with existing
+            for exp_id, row in existing.items():
+                merged_map[exp_id] = row
+            
+            # Apply incoming
             for exp in experiments:
-                if self._validate_experiment_structure(exp):
-                    # Ensure last_updated is set
-                    if 'last_updated' not in exp:
-                        exp['last_updated'] = datetime.now().isoformat()
-                    valid_experiments.append(exp)
-                else:
+                if not self._validate_experiment_structure(exp):
                     logger.error(f"❌ Invalid experiment structure: {exp.get('experiment_id', 'unknown')}")
                     return False
-            
-            # Create dataset
-            dataset = Dataset.from_list(valid_experiments)
+                exp_id = exp['experiment_id']
+                incoming = exp
+                if exp_id not in merged_map:
+                    incoming['last_updated'] = incoming.get('last_updated') or datetime.now().isoformat()
+                    merged_map[exp_id] = incoming
+                    continue
+                # Merge with existing
+                base = merged_map[exp_id]
+                # Parse JSON fields
+                base_metrics = _parse_json_field(base.get('metrics'), [])
+                base_params = _parse_json_field(base.get('parameters'), {})
+                base_artifacts = _parse_json_field(base.get('artifacts'), [])
+                base_logs = _parse_json_field(base.get('logs'), [])
+                inc_metrics = _parse_json_field(incoming.get('metrics'), [])
+                inc_params = _parse_json_field(incoming.get('parameters'), {})
+                inc_artifacts = _parse_json_field(incoming.get('artifacts'), [])
+                inc_logs = _parse_json_field(incoming.get('logs'), [])
+                # Merge metrics with de-dup
+                merged_metrics = []
+                seen = set()
+                for entry in base_metrics + inc_metrics:
+                    try:
+                        # Use the original entry so _metrics_key can properly
+                        # distinguish dict vs non-dict entries
+                        key = _metrics_key(entry)
+                    except Exception:
+                        key = (None, None)
+                    if key not in seen:
+                        seen.add(key)
+                        merged_metrics.append(entry)
+                # Merge params
+                merged_params = {}
+                if isinstance(base_params, dict):
+                    merged_params.update(base_params)
+                if isinstance(inc_params, dict):
+                    merged_params.update(inc_params)
+                # Merge artifacts and logs with de-dup
+                def _dedup_list(lst):
+                    out = []
+                    seen_local = set()
+                    for item in lst:
+                        key = json.dumps(item, sort_keys=True, default=str) if not isinstance(item, str) else item
+                        if key not in seen_local:
+                            seen_local.add(key)
+                            out.append(item)
+                    return out
+                merged_artifacts = _dedup_list(list(base_artifacts) + list(inc_artifacts))
+                merged_logs = _dedup_list(list(base_logs) + list(inc_logs))
+                # Rebuild merged record preferring incoming scalars
+                merged_rec = dict(base)
+                merged_rec.update({k: v for k, v in incoming.items() if k not in ('metrics', 'parameters', 'artifacts', 'logs')})
+                merged_rec['metrics'] = json.dumps(merged_metrics, default=str)
+                merged_rec['parameters'] = json.dumps(merged_params, default=str)
+                merged_rec['artifacts'] = json.dumps(merged_artifacts, default=str)
+                merged_rec['logs'] = json.dumps(merged_logs, default=str)
+                merged_rec['last_updated'] = datetime.now().isoformat()
+                merged_map[exp_id] = merged_rec
+            
+            # Prepare final list
+            valid_experiments = list(merged_map.values())
+            # Ensure all have mandatory fields encoded
+            normalized = []
+            for rec in valid_experiments:
+                # Normalize json fields to strings
+                for f, default in (('metrics', []), ('parameters', {}), ('artifacts', []), ('logs', [])):
+                    val = rec.get(f)
+                    if not isinstance(val, str):
+                        rec[f] = json.dumps(val if val is not None else default, default=str)
+                if 'last_updated' not in rec:
+                    rec['last_updated'] = datetime.now().isoformat()
+                normalized.append(rec)
+            
+            dataset = Dataset.from_list(normalized)
             
             # Generate commit message if not provided
             if not commit_message:
-                commit_message = f"Update dataset with {len(valid_experiments)} experiments ({datetime.now().isoformat()})"
+                commit_message = f"Union-merge update with {len(normalized)} experiments ({datetime.now().isoformat()})"
             
             # Push to hub
             dataset.push_to_hub(
@@ -163,7 +267,7 @@ class TrackioDatasetManager:
                 commit_message=commit_message
             )
             
-            logger.info(f"✅ Successfully saved {len(valid_experiments)} experiments to {self.dataset_repo}")
+            logger.info(f"✅ Successfully saved {len(normalized)} experiments (union-merged) to {self.dataset_repo}")
             return True
             
         except Exception as e: