File size: 8,343 Bytes
3270dae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | """Async batch iterator for training with background tokenization."""
from typing import Dict, List, Optional, Any, Iterator
import torch
from taoTrain.data.tokenization_queue import TokenizationQueue
from taoTrain.data.sft_utils import build_response_only_next_token_labels
class AsyncBatchIterator:
"""
Iterator that yields batches from a tokenization queue.
This allows batches to be consumed directly from the background tokenization
thread without waiting for all chunks to be tokenized upfront.
The iterator:
1. Pulls pre-tokenized chunks from the TokenizationQueue
2. Yields individual samples or batches
3. Handles movement to device (GPU/CPU) at batch level
4. Supports gradient accumulation
"""
def __init__(
self,
tokenization_queue: TokenizationQueue,
batch_size: int,
device: torch.device,
drop_last: bool = True,
gradient_accumulation_steps: int = 1,
):
"""
Initialize async batch iterator.
Args:
tokenization_queue: TokenizationQueue instance
batch_size: Batch size for yielding batches
device: torch.device to move batches to
drop_last: If True, drop last incomplete batch
gradient_accumulation_steps: For logging purposes (not used here)
"""
self.queue = tokenization_queue
self.batch_size = batch_size
self.device = device
self.drop_last = drop_last
self.gradient_accumulation_steps = gradient_accumulation_steps
# State for iteration
self._current_chunk: Optional[Dict[str, List]] = None
self._current_idx = 0
self._samples_yielded = 0
self._finished = False
def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
"""Return iterator (self)."""
# Reset state for new epoch
self._current_chunk = None
self._current_idx = 0
self._samples_yielded = 0
self._finished = False
# Reset tokenization queue for epochs 2+
if self.queue._next_chunk_idx > 0:
print(f"\n✓ Resetting TokenizationQueue for next epoch (cur_idx={self.queue._next_chunk_idx})")
self.queue.reset_for_next_epoch()
# Start tokenization threads once per iterator creation
if not self.queue._threads:
print("\n✓ Starting TokenizationQueue worker threads...")
self.queue.start()
else:
print(f"\n⚠ TokenizationQueue threads already running: {len(self.queue._threads)} active")
return self
def __next__(self) -> Dict[str, torch.Tensor]:
"""
Get next batch.
Yields:
Dict with 'input_ids', 'attention_mask', 'labels' (all as torch tensors on device)
Raises:
StopIteration: When no more batches available
"""
batch = self._get_next_batch()
if batch is None:
print("AsyncBatchIterator: No more batches available, stopping iteration.")
raise StopIteration
return batch
def _get_next_batch(self) -> Optional[Dict[str, torch.Tensor]]:
"""
Fetch and collate the next batch.
Returns:
Dict with batch tensors, or None if iteration exhausted
"""
batch_input_ids = []
batch_attention_masks = []
batch_labels = []
while len(batch_input_ids) < self.batch_size:
# Try to get next sample from current chunk
if self._current_chunk is None or self._current_idx >= len(self._current_chunk["input_ids"]):
# Need new chunk
self._current_chunk = self.queue.get_next_chunk(timeout=30.0) # 30s polling timeout
if self._current_chunk is None:
if not self.queue.is_exhausted:
continue
# Queue exhausted
chunk_count = self.queue._next_chunk_idx if hasattr(self.queue, '_next_chunk_idx') else 'unknown'
print(f"AsyncBatchIterator: No more chunks (processed {chunk_count}/{len(self.queue._chunk_order)})")
print(f"AsyncBatchIterator: Samples yielded so far: {self._samples_yielded}")
self._finished = True
break
self._current_idx = 0
# Get sample from current chunk
input_ids = self._current_chunk["input_ids"][self._current_idx]
attention_mask = self._current_chunk["attention_mask"][self._current_idx]
# Generate labels based on SFT or pretrain mode
if "mask" in self._current_chunk:
# SFT mode: use mask to determine which tokens to train on
# mask=0 → label=-100 (ignore), mask=1 → label=input_id (train on)
mask = self._current_chunk["mask"][self._current_idx]
labels = build_response_only_next_token_labels(input_ids, mask)
else:
# Pretrain mode: shift labels by 1 for next-token prediction
# Position i predicts token at position i+1
labels = input_ids[1:] + [-100] # Append -100 as final position
# Mark padding tokens as -100 to ignore in loss computation
for i, mask_val in enumerate(attention_mask):
if mask_val == 0:
labels[i] = -100
batch_input_ids.append(input_ids)
batch_attention_masks.append(attention_mask)
batch_labels.append(labels)
self._current_idx += 1
self._samples_yielded += 1
# Return batch if we have any samples, respecting drop_last
if len(batch_input_ids) == 0:
print(f"AsyncBatchIterator: No samples collected for batch. Finished={self._finished}, returning None.")
return None
if len(batch_input_ids) < self.batch_size and self.drop_last:
incomplete_pct = (len(batch_input_ids) / self.batch_size) * 100
print(f"AsyncBatchIterator: Batch incomplete ({len(batch_input_ids)}/{self.batch_size} = {incomplete_pct:.1f}%) and drop_last=True, returning None.")
return None
return self._collate_batch(batch_input_ids, batch_attention_masks, batch_labels)
def _collate_batch(
self,
batch_input_ids: List[List[int]],
batch_attention_masks: List[List[int]],
batch_labels: List[List[int]],
) -> Dict[str, torch.Tensor]:
"""
Collate batch samples and move to device.
Args:
batch_input_ids: List of token ID lists
batch_attention_masks: List of attention mask lists
batch_labels: List of label lists
Returns:
Collated batch as torch tensors on device
"""
# Convert to tensors
input_ids_tensor = torch.tensor(batch_input_ids, dtype=torch.long, device=self.device)
attention_mask_tensor = torch.tensor(batch_attention_masks, dtype=torch.long, device=self.device)
labels_tensor = torch.tensor(batch_labels, dtype=torch.long, device=self.device)
return {
"input_ids": input_ids_tensor,
"attention_mask": attention_mask_tensor,
"labels": labels_tensor,
}
def __len__(self) -> int:
"""Return approximate number of batches."""
total_samples = len(self.queue)
if self.drop_last:
return total_samples // self.batch_size
else:
return (total_samples + self.batch_size - 1) // self.batch_size
def shutdown(self):
"""Shutdown the async iterator and background thread."""
self.queue.shutdown(wait=True)
def __del__(self):
"""Cleanup on deletion."""
self.shutdown()
|