adds flash attention 3 kernel

config/train_gpt_oss_openhermes_fr_memory_optimized.py

@@ -35,7 +35,7 @@ config = GPTOSSEnhancedCustomConfig(
     # Dataset sampling optimized for memory constraints
     max_samples=800000,                      # Reduced from 800K for memory efficiency
     min_length=15,                          # Slightly higher minimum for quality
-    max_length=2048,                        # Explicit max length for memory control
+    max_length=4096,                        # Explicit max length for memory control
     
     # ============================================================================
     # MEMORY-OPTIMIZED TRAINING HYPERPARAMETERS
@@ -56,7 +56,7 @@ config = GPTOSSEnhancedCustomConfig(
     # MODEL CONFIGURATION - Memory Optimized for GPT-OSS
     # ============================================================================
     model_name="openai/gpt-oss-20b",
-    max_seq_length=1024,                    # Reduced from 3072 for memory optimization
+    max_seq_length=4096,                    # Reduced from 3072 for memory optimization
     use_flash_attention=True,               # Critical for memory efficiency
     use_gradient_checkpointing=True,        # Essential for memory optimization
     
@@ -106,7 +106,7 @@ config = GPTOSSEnhancedCustomConfig(
     # ============================================================================
     # Model loading with memory constraints
     model_kwargs={
-        "attn_implementation": "eager",     # Memory-safe attention
+        "attn_implementation": "kernels-community/vllm-flash-attn3",  # Much faster attention on A100/H100
         "torch_dtype": "auto",              # Let model decide (MXFP4 compatible)
         "use_cache": False,                 # Disable KV cache for training
         "device_map": "auto",               # Automatic device mapping
@@ -114,10 +114,10 @@ config = GPTOSSEnhancedCustomConfig(
         "max_memory": {0: "75GB"},          # Reserve memory for other processes
     },
     
-    # Data loading optimized for memory efficiency
-    dataloader_num_workers=2,               # Reduced workers to save memory
-    dataloader_pin_memory=False,            # Disable to save memory
-    dataloader_prefetch_factor=1,           # Minimal prefetch for memory
+    # Data loading optimized for throughput
+    dataloader_num_workers=4,                # More workers for faster loading
+    dataloader_pin_memory=True,              # Pin memory for faster host->GPU copies
+    dataloader_prefetch_factor=2,            
     
     # Memory management optimizations
     max_memory_per_gpu="75GB",              # Explicit memory limit
@@ -126,7 +126,7 @@ config = GPTOSSEnhancedCustomConfig(
     remove_unused_columns=True,             # Remove unnecessary data
     
     # ============================================================================
-    # EVALUATION & LOGGING - Memory Efficient
+    # EVALUATION & LOGGING - Fast Iterations
     # ============================================================================
     eval_strategy="steps",
     eval_steps=500,                         # Less frequent evaluation for memory
@@ -134,7 +134,7 @@ config = GPTOSSEnhancedCustomConfig(
     
     save_strategy="steps", 
     save_steps=1000,                        # Less frequent saves for memory/storage
-    save_total_limit=2,                     # Keep only 2 checkpoints for memory
+    save_total_limit=3,                     # Keep only 2 checkpoints for memory
     save_only_model=True,                   # Save only model weights
     
     metric_for_best_model="eval_loss",

requirements/requirements_core.txt

@@ -5,6 +5,7 @@ datasets>=2.14.0
 accelerate>=0.20.0
 peft>=0.17.0  # Updated for GPT-OSS LoRA support
 trl>=0.20.0  # Updated for GPT-OSS compatibility
+kernels
 
 # Hugging Face Hub for model and space management
 huggingface_hub>=0.19.0

scripts/trackio_tonic/app.py

@@ -25,10 +25,11 @@ class TrackioSpace:
     def __init__(self, hf_token: Optional[str] = None, dataset_repo: Optional[str] = None):
         self.experiments = {}
         self.current_experiment = None
+        self.backup_mode = False
         
         # Get dataset repository and HF token from parameters or environment variables
-        # Use dynamic default based on environment or fallback to generic default
-        default_dataset_repo = os.environ.get('TRACKIO_DATASET_REPO', 'trackio-experiments')
+        # Respect explicit values; avoid hardcoded defaults that might point to test repos
+        default_dataset_repo = os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
         self.dataset_repo = dataset_repo or default_dataset_repo
         self.hf_token = hf_token or os.environ.get('HF_TOKEN')
         
@@ -75,12 +76,14 @@ class TrackioSpace:
                     # Fall back to backup data
                     self._load_backup_experiments()
             else:
-                # No HF token, use backup data
+                # No HF token, use backup data but do not allow saving to dataset from backup
                 self._load_backup_experiments()
+                self.backup_mode = True
                 
         except Exception as e:
             logger.error(f"Failed to load experiments: {e}")
             self._load_backup_experiments()
+            self.backup_mode = True
     
     def _load_backup_experiments(self):
         """Load backup experiments when dataset is not available"""
@@ -314,6 +317,9 @@ class TrackioSpace:
     def _save_experiments(self):
         """Save experiments to HF Dataset"""
         try:
+            if self.backup_mode:
+                logger.warning("⚠️ Backup mode active; skipping dataset save to avoid overwriting real data with demo values")
+                return
             if self.hf_token:
                 from datasets import Dataset
                 from huggingface_hub import HfApi
@@ -565,17 +571,20 @@ def create_dataset_repository(hf_token: str, dataset_repo: str) -> str:
     except Exception as e:
         return f"❌ Failed to create dataset: {str(e)}\n\n💡 Troubleshooting:\n1. Check your HF token has write permissions\n2. Verify the username in the repository name\n3. Ensure the dataset name is valid"
 
-# Initialize API client for remote data
+# Initialize API client for remote data if environment provides a space id/url
 api_client = None
 try:
-    from trackio_api_client import create_trackio_client
-    api_client = create_trackio_client()
-    if api_client:
+    from trackio_api_client import TrackioAPIClient
+    space_id = os.environ.get('TRACKIO_URL') or os.environ.get('TRACKIO_SPACE_ID')
+    if space_id:
+        api_client = TrackioAPIClient(space_id, os.environ.get('HF_TOKEN'))
         logger.info("✅ API client initialized for remote data access")
     else:
-        logger.warning("⚠️ Could not initialize API client, using local data only")
+        logger.info("No TRACKIO_URL/TRACKIO_SPACE_ID set; remote API client disabled")
 except ImportError:
     logger.warning("⚠️ API client not available, using local data only")
+except Exception as e:
+    logger.warning(f"⚠️ Could not initialize API client: {e}")
 
 # Add Hugging Face Spaces compatibility
 def is_huggingface_spaces():
@@ -590,8 +599,8 @@ def get_persistent_data_path():
     else:
         return "trackio_experiments.json"
 
-# Override the data file path for HF Spaces
-if is_huggingface_spaces():
+# Override the data file path for HF Spaces if attribute exists
+if is_huggingface_spaces() and hasattr(trackio_space, 'data_file'):
     logger.info("🚀 Running on Hugging Face Spaces - using persistent storage")
     trackio_space.data_file = get_persistent_data_path()
 

scripts/training/train_gpt_oss.py

@@ -458,11 +458,17 @@ def split_dataset(dataset, config):
 def setup_trackio_tracking(config):
     """Setup Trackio tracking if enabled"""
     
-    if not config.enable_tracking or not config.trackio_url:
+    if not getattr(config, 'enable_tracking', False):
         print("Trackio tracking disabled or URL not provided")
         return None
     
-    print(f"Setting up Trackio tracking: {config.trackio_url}")
+    # Resolve Trackio URL from config or environment
+    trackio_url = getattr(config, 'trackio_url', None) or os.environ.get('TRACKIO_URL') or os.environ.get('TRACKIO_SPACE_ID')
+    if not trackio_url:
+        print("Trackio tracking enabled but no TRACKIO_URL/TRACKIO_SPACE_ID provided; skipping Trackio setup")
+        return None
+
+    print(f"Setting up Trackio tracking: {trackio_url}")
     
     # Import the correct TrackioAPIClient
     import sys
@@ -472,8 +478,8 @@ def setup_trackio_tracking(config):
     
     # Initialize Trackio client using the correct API
     trackio_client = TrackioAPIClient(
-        space_id=config.trackio_url,
-        hf_token=config.trackio_token
+        space_id=trackio_url,
+        hf_token=getattr(config, 'trackio_token', None) or os.environ.get('HF_TOKEN')
     )
     
     return trackio_client

src/monitoring.py

@@ -120,12 +120,18 @@ class SmolLM3Monitor:
         """Setup Trackio API client"""
         try:
             # Get Trackio configuration from environment or parameters
-            space_id = trackio_url or os.getenv('TRACKIO_SPACE_ID')
-            
+            # Accept either a full URL or an org/space identifier
+            # Prefer explicit parameter, then environment variables
+            space_id = (
+                trackio_url
+                or os.getenv('TRACKIO_URL')
+                or os.getenv('TRACKIO_SPACE_ID')
+            )
+
             if not space_id:
-                # Use the deployed Trackio Space ID
-                space_id = "Tonic/trackio-monitoring-20250727"
-                logger.info(f"Using default Trackio Space ID: {space_id}")
+                logger.warning("No Trackio Space configured via param or env (TRACKIO_URL/TRACKIO_SPACE_ID). Disabling Trackio tracking.")
+                self.enable_tracking = False
+                return
             
             # Get HF token for Space resolution
             hf_token = self.hf_token or trackio_token or os.getenv('HF_TOKEN')

src/trackio.py

@@ -40,7 +40,12 @@ def init(
             project_name = os.environ.get('EXPERIMENT_NAME', 'smollm3_experiment')
         
         # Extract configuration from kwargs
-        trackio_url = kwargs.get('trackio_url') or os.environ.get('TRACKIO_URL')
+        # Accept both TRACKIO_URL (full URL or org/space) and TRACKIO_SPACE_ID
+        trackio_url = (
+            kwargs.get('trackio_url')
+            or os.environ.get('TRACKIO_URL')
+            or os.environ.get('TRACKIO_SPACE_ID')
+        )
         trackio_token = kwargs.get('trackio_token') or os.environ.get('TRACKIO_TOKEN')
         hf_token = kwargs.get('hf_token') or os.environ.get('HF_TOKEN')
         dataset_repo = kwargs.get('dataset_repo') or os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')

templates/spaces/trackio/app.py

@@ -27,6 +27,7 @@ class TrackioSpace:
     def __init__(self, hf_token: Optional[str] = None, dataset_repo: Optional[str] = None):
         self.experiments = {}
         self.current_experiment = None
+        self.using_backup_data = False
         
         # Get dataset repository and HF token from parameters or environment variables
         self.dataset_repo = dataset_repo or os.environ.get('TRACKIO_DATASET_REPO', 'Tonic/trackio-experiments')
@@ -80,10 +81,11 @@ class TrackioSpace:
                     reverse=True
                 ))
                 
-                # If no experiments found, use backup
+                # If no experiments found, use backup but mark backup mode to avoid accidental writes
                 if not self.experiments:
                     logger.info("📊 No experiments found in dataset, using backup data")
                     self._load_backup_experiments()
+                    self.using_backup_data = True
                 
                 return
             
@@ -91,15 +93,18 @@ class TrackioSpace:
             if self.hf_token:
                 success = self._load_experiments_direct()
                 if success:
+                    self.using_backup_data = False
                     return
             
             # Final fallback to backup data
             logger.info("🔄 Using backup data")
             self._load_backup_experiments()
+            self.using_backup_data = True
                 
         except Exception as e:
             logger.error(f"❌ Failed to load experiments: {e}")
             self._load_backup_experiments()
+            self.using_backup_data = True
     
     def _load_experiments_direct(self) -> bool:
         """Load experiments directly from HF Dataset without dataset manager"""
@@ -423,6 +428,9 @@ class TrackioSpace:
     def _save_experiments(self):
         """Save experiments to HF Dataset with data preservation"""
         try:
+            if self.using_backup_data:
+                logger.warning("⚠️ Using backup data; skip saving to dataset to avoid overwriting with demo values")
+                return
             # Use dataset manager for safe operations if available
             if self.dataset_manager:
                 logger.info("💾 Saving experiments using dataset manager (data preservation)")
@@ -782,21 +790,27 @@ def create_dataset_repository(hf_token: str, dataset_repo: str) -> str:
     except Exception as e:
         return f"❌ Failed to create dataset: {str(e)}\n\n💡 Troubleshooting:\n1. Check your HF token has write permissions\n2. Verify the username in the repository name\n3. Ensure the dataset name is valid\n4. Check internet connectivity"
 
-# Initialize API client for remote data
+"""
+Initialize API client for remote data. We do not hardcode a default test URL to avoid
+overwriting dataset content with demo data. The API client will only be initialized
+when TRACKIO_URL or TRACKIO_SPACE_ID is present.
+"""
 api_client = None
 try:
     from trackio_api_client import TrackioAPIClient
-    # Get Trackio URL from environment or use default
-    trackio_url = os.environ.get('TRACKIO_URL', 'https://tonic-test-trackio-test.hf.space')
-    
-    # Clean up URL to avoid double protocol issues
-    if trackio_url.startswith('https://https://'):
-        trackio_url = trackio_url.replace('https://https://', 'https://')
-    elif trackio_url.startswith('http://http://'):
-        trackio_url = trackio_url.replace('http://http://', 'http://')
-    
-    api_client = TrackioAPIClient(trackio_url)
-    logger.info(f"✅ API client initialized for remote data access: {trackio_url}")
+    # Resolve Trackio space from environment
+    trackio_url_env = os.environ.get('TRACKIO_URL') or os.environ.get('TRACKIO_SPACE_ID')
+    if trackio_url_env:
+        # Clean up URL to avoid double protocol issues
+        trackio_url = trackio_url_env
+        if trackio_url.startswith('https://https://'):
+            trackio_url = trackio_url.replace('https://https://', 'https://')
+        elif trackio_url.startswith('http://http://'):
+            trackio_url = trackio_url.replace('http://http://', 'http://')
+        api_client = TrackioAPIClient(trackio_url)
+        logger.info(f"✅ API client initialized for remote data access: {trackio_url}")
+    else:
+        logger.info("No TRACKIO_URL/TRACKIO_SPACE_ID set; remote API client disabled")
 except ImportError:
     logger.warning("⚠️ API client not available, using local data only")
 except Exception as e: