accelerate-streaming-dataset-loading

[david-thrower]

Title: Streaming Dataset Performance Bottleneck and Progress Bar Issues

Description

The API enhancements for handling streaming datasets implemented in #34 have
two critical issues affecting training time and monitoring accuracy.

Issue 1: CPU Bottleneck in Preprocessing

Expected vs. Actual Performance:

• Config: 512 seq len, d_model 512, 1.5B tokens, 3 recurrent loops
• Throughput: 29,000 tokens/sec (healthy for 3 loops)
• Expected completion: 9-12 hours
• Actual: >18 hours with no epoch completion

Performance Calculation:

Factor Impact
4x data (1.5B vs 400M token baseline) 4x train time
4x batching time +4x overhead
1/4 steps -4x time (cumulatively net 0 change)
2x grad accumulation, 0.5x batch size net 0 (cumulatively net 0 change)
1.5x recurrence 1.5x time (cumulatively net 1.5 X change)
Linear attention (4x seq_len benefit) ~4x time (cumulatively net 4.5 change)
Net expected ~4.5x the 2 1/2 hour baseline ~9-12 hours

Suspected Cause: CPU-bound batched preprocessing (tokenization) creating GPU
idle time. Previous 400M token run with 2 loops completed in ~2.5 hours using
List[str] instead of iterable dataset.

Questions:

1 Is the bottleneck confirmation correct, or is there a math error?
2 If bottleneck is real, what are the options to satisfy both:
• Streaming (don't load everything in memory)
• No batching wait bottleneck

Potential Solutions:

• Pre-tokenize and cache to disk (count steps upfront)
• Stream on-the-fly with preprocessing bottleneck fixed

Issue 2: Progress Bar Regression

Current Behavior (Streaming):

Epoch 1: 256527 batch [18:13:53, 3.91batch/s, loss=5.1309, ppl=169.16,
lr=1.00e-03, tok/s=29,007]

• Missing percent completion
• Batch number shown without total batches
• Estimated batch count exists but TQDM disabled

Expected Behavior (List[str]):

Epoch 25: 100%|██████████| 101/101 [01:54<00:00, 1.13s/batch, loss=2.4469,
ppl=11.55, lr=2.63e-04, tok/s=679]

• Shows percent completion
• Shows current/total batches
• Full progress bar

Root Cause: Streaming datasets were set up to use estimated steps, and TQDM
was disabled. The corrected logic to count batches first exists but may not be
properly integrated.

[david-thrower]

Summary of [agent proposed] Changes in streaming-data-tests-copy/HelixLM Branch

Files Modified

• helix_lm/dataset.py
• helix_lm/trainer.py

---

Analysis: What Was Changed

Issue 1: CPU Bottleneck (Addressed via New Dataset Architecture)

Root Cause: The original implementation in DocumentAwareDataset.len()
(lines 316-318 in original) forced eager chunking when len() was called:

   def __len__(self):                                                           
       self._ensure_chunks()  # <-- This builds ALL chunks just to count them!  
       return len(self.chunks)                                                  

When Trainer.train_epoch() called len(self.train_loader) to set up the
scheduler, this triggered complete tokenization and chunking of the entire
dataset, causing the massive startup delay observed (counting 256K+ batches).

Changes Made:

1 Three new dataset strategies introduced in dataset.py:

Strategy Class Use Case
────────────────────────────────────────────────────────────────────────────
Map-style (existing) DocumentAwareDataset, Datasets that fit in
HelixDataset memory
Streaming iterable (NEW) HelixIterableDataset 2B+ token corpora, true
O(1) memory
Pre-chunked (NEW) HelixPrechunkedDataset Fast repeatable training
with pre-tokenized data

2 HelixIterableDataset (lines 502-731 NEW):
• True streaming with per-document on-the-fly tokenization
• No len() method - raises TypeError to signal "iterable dataset"
• Reservoir shuffle with configurable buffer size
• set_epoch() for reshuffling between epochs

3 HelixPrechunkedDataset (lines 738-951 NEW):
• Pre-chunks entire dataset once using Dataset.map(batched=True,
num_proc=...)
• Saves to disk as Arrow-backed dataset
• Fast random access for training

4 Unified factory function create_unified_data_loader() (lines 1147-1275
NEW):
• Auto-detects List[str] vs iterable/streaming
• Routes to appropriate dataset implementation

---

Issue 2: Progress Bar (Addressed via Iterable Detection + Counting Pass)

Root Cause: Original Trainer assumed all datasets have len(). For iterable
datasets, TQDM showed batch [HH:MM:SS, ...] format without percentage.

Changes in trainer.py:

   1 New helper _is_iterable_dataset() (lines 93-104):
                                                                                
      def _is_iterable_dataset(dataset):                                        
          """Check if a dataset is iterable (no random access / len())."""      
          if hasattr(dataset, "__iter__") and not hasattr(dataset,              
      "__getitem__"):                                                           
              return True                                                       
          # Also catches HelixIterableDataset which has __getitem__ but         
      TypeError on __len__                                                      

2 Lazy scheduler initialization _init_scheduler() (lines 345-400):
• For iterable datasets: performs counting pass first via
_count_iterable_dataset()
• Then resets dataset with set_epoch(0) for actual training
• Uses cached count for accurate TQDM total

3 Progress bar with total= for iterable datasets (lines 422-433):

      # For iterable datasets with known batch count, enable progress bar with  
      total                                                                     
      pbar_total = None                                                         
      if self._is_iterable and self._cached_dataset_length is not None:         
          pbar_total = self._cached_dataset_length                              
                                                                                
      pbar = tqdm(                                                              
          self.train_loader,                                                    
          desc=f"Epoch {epoch}",                                                
          unit="batch",                                                         
          disable=not self.verbose,                                             
          total=pbar_total,  # <-- Now shows percentage!                        
      )                                                                         

4 warmup_ratio parameter (lines 133, 169, 362-375):
• Alternative to warmup_steps for cases where exact step count is unknown
• warmup_ratio=0.1 = 10% of total steps are warmup

---

Summary Table: Before vs After

                                              After                         

Aspect Before (Issue #34) (streaming-data-tests-copy)
────────────────────────────────────────────────────────────────────────────
Dataset types Only map-style (List[str]) Map + Iterable + Pre-chunked
Memory scaling O(n) - must fit entire O(1) streaming or cached
dataset pre-chunked
Startup delay Forced full chunking on Delayed until first iteration
len() or cached
Scheduler setup Required exact step count up Supports warmup_ratio or
front counts iterables
Progress bar No percentage for unknown Counting pass enables full
lengths TQDM with %
Shuffle DataLoader shuffle (map Reservoir shuffle for
only) iterables

---

Recommendation for Issue

These changes address both Issue 1 (CPU bottleneck from eager len() calls) and
Issue 2 (progress bar without percentage). The key insight is that the
original DocumentAwareDataset.len() implementation was the bottleneck -
calling it forced complete tokenization.

The new architecture provides three legitimate paths:

1 Small datasets (<500M tokens): Continue using List[str] →
DocumentAwareDataset
2 Large streaming (2B+ tokens): Use HelixIterableDataset with counting pass
for accurate progress
3 Fast repeatable training: Use HelixPrechunkedDataset with pre-tokenized
shards