File size: 16,865 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 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 | """Background tokenization queue for streaming large JSONL datasets."""
import queue
import threading
import time
from typing import Dict, List, Optional, Any, Callable
import torch
from taoTrain.data.chunk_manager import ChunkManager
class TokenizationQueue:
"""
Background threads that continuously tokenize chunks and stores them in a queue.
This allows tokenization to happen in parallel with training, avoiding the bottleneck
of tokenizing all data upfront before training starts.
Supports multiple worker threads for faster throughput. Each thread greedily
grabs the next available chunk using an atomic counter.
Attributes:
total_items: Total number of samples across all chunks
queue_size: Maximum number of chunks to buffer in memory
num_threads: Number of worker threads for tokenization
"""
def __init__(
self,
chunk_manager: ChunkManager,
tokenizer: Any,
config: "TrainingConfig", # type: ignore
max_queue_size: int = 2,
shuffle_chunks: bool = True,
num_threads: int = 1,
):
"""
Initialize tokenization queue with multithreading support.
Args:
chunk_manager: ChunkManager instance loaded with chunks
tokenizer: Tokenizer instance (HuggingFace or SentencePiece wrapper)
config: Training configuration with model and dataset settings
max_queue_size: Maximum chunks to buffer in queue (memory constraint)
shuffle_chunks: Whether to shuffle chunk order at initialization
num_threads: Number of worker threads for tokenization (default: 1)
Raises:
ValueError: If chunk_manager has no chunks or num_threads < 1
"""
if chunk_manager.num_chunks == 0:
raise ValueError("ChunkManager must have at least one chunk")
if num_threads < 1:
raise ValueError(f"num_threads must be >= 1, got {num_threads}")
self.chunk_manager = chunk_manager
self.tokenizer = tokenizer
self.config = config
self.max_queue_size = max_queue_size
self.shuffle_chunks = shuffle_chunks
self.num_threads = num_threads
# Detect SFT mode: check for response_loss_only flag
self.is_sft_mode = hasattr(config, 'response_loss_only') and config.response_loss_only
# Calculate total items across all chunks
self.total_items = chunk_manager.effective_lines
# Thread-safe queue for tokenized chunks
self._queue: queue.Queue[Dict[str, List]] = queue.Queue(maxsize=max_queue_size)
# Control signals
self._stop_event = threading.Event()
self._error_event = threading.Event()
self._error_messages: List[str] = []
self._threads: List[threading.Thread] = []
# Thread-safe chunk distribution
self._next_chunk_idx = 0
self._chunk_idx_lock = threading.Lock()
self._active_threads = 0
self._active_threads_lock = threading.Lock()
# Chunk ordering
self._chunk_order = list(range(chunk_manager.num_chunks))
print(f"TokenizationQueue initialized with {chunk_manager.num_chunks} chunks, total {chunk_manager.effective_lines} samples")
print(f"Using {num_threads} tokenization worker thread{'s' if num_threads != 1 else ''}")
print(f"Max queue size: {max_queue_size} chunks (memory constraint)")
if self.shuffle_chunks:
import random
random.shuffle(self._chunk_order)
def _get_next_chunk_idx(self) -> Optional[int]:
"""
Atomically get the next chunk index for processing.
Returns:
Chunk index to process, or None if all chunks have been assigned
"""
with self._chunk_idx_lock:
if self._next_chunk_idx < len(self._chunk_order):
chunk_idx = self._chunk_order[self._next_chunk_idx]
self._next_chunk_idx += 1
return chunk_idx
return None
def start(self):
"""Start the tokenization background worker threads."""
if self._threads:
raise RuntimeError(f"Tokenization threads already started ({len(self._threads)} active)")
# Create and start N worker threads
for thread_id in range(self.num_threads):
thread = threading.Thread(target=self._worker, args=(thread_id,), daemon=False)
self._threads.append(thread)
thread.start()
def _worker(self, thread_id: int):
"""
Worker thread target: greedy chunk processing with thread-safe distribution.
Args:
thread_id: Identifier for this worker thread
"""
with self._active_threads_lock:
self._active_threads += 1
try:
while True:
# Check for stop signal
if self._stop_event.is_set():
break
# Get next chunk to process (atomic operation)
chunk_num = self._get_next_chunk_idx()
if chunk_num is None:
# All chunks assigned
break
# Load chunk
chunk_examples = self.chunk_manager.read_chunk(chunk_num)
# Tokenize chunk based on mode
if self.is_sft_mode:
tokenized_chunk = self._tokenize_batch_sft(chunk_examples)
else:
# Extract texts for pretrain
text_field = self.config.dataset.text_field
texts = [obj.get(text_field, "") for obj in chunk_examples]
tokenized_chunk = self._tokenize_batch(texts)
# Put in queue (blocks if queue is full)
self._queue.put(tokenized_chunk)
print(f"[Worker-{thread_id}] Processed chunk {chunk_num}, put {len(tokenized_chunk['input_ids'])} samples in queue")
except Exception as e:
error_msg = f"[Worker-{thread_id}] {str(e)}"
print(f"Worker-{thread_id} encountered an error: {error_msg}")
# Thread-safe append to error list
self._error_messages.append(error_msg)
self._error_event.set()
finally:
with self._active_threads_lock:
self._active_threads -= 1
remaining = self._active_threads
print(f"[Worker-{thread_id}] Finished processing. Active threads remaining: {remaining}")
def _tokenize_batch(self, texts: List[str]) -> Dict[str, List]:
"""
Tokenize a batch of texts, join with EOS, and split into fixed-size sequences.
This packs multiple documents into longer sequences separated by EOS tokens,
then splits the concatenated tokens into N fixed-size chunks of max_seq_length.
Args:
texts: List of text strings
Returns:
Dict with 'input_ids' and 'attention_mask' lists, where each element
is a fixed-size sequence of length max_seq_length
"""
max_seq_length = self.config.model.max_seq_length
# Get EOS token ID
eos_token_id = self.tokenizer.eos_token_id
unk_token_id = self.tokenizer.unk_token_id
if eos_token_id is None:
raise ValueError("Tokenizer does not have an EOS token defined")
if unk_token_id is None:
raise ValueError("Tokenizer does not have an UNK token defined")
# Tokenize all texts without truncation
all_token_ids = []
for i, text in enumerate(texts):
tokenized = self.tokenizer(
text,
truncation=False,
return_attention_mask=False,
)
# Remove UNK tokens from tokenized output (if any)
tokenized["input_ids"] = [tid for tid in tokenized["input_ids"] if tid != unk_token_id]
all_token_ids.extend(tokenized["input_ids"])
# Add EOS token between documents (except after the last one)
if i < len(texts) - 1:
all_token_ids.append(eos_token_id)
# Split into N fixed-size sequences
sequences_input_ids = []
sequences_attention_masks = []
for i in range(0, len(all_token_ids), max_seq_length):
seq = all_token_ids[i : i + max_seq_length]
# Pad sequence if it's shorter than max_seq_length
if len(seq) < max_seq_length:
# Create attention mask before padding
attention_mask = [1] * len(seq) + [0] * (max_seq_length - len(seq))
# Pad with 0 (assuming 0 is the pad token, or use tokenizer.pad_token_id)
pad_token_id = self.tokenizer.pad_token_id or 0
seq = seq + [pad_token_id] * (max_seq_length - len(seq))
else:
attention_mask = [1] * max_seq_length
sequences_input_ids.append(seq)
sequences_attention_masks.append(attention_mask)
return {
"input_ids": sequences_input_ids,
"attention_mask": sequences_attention_masks,
}
def _tokenize_batch_sft(self, records: List[Dict[str, Any]]) -> Dict[str, List]:
"""
Tokenize a batch of SFT records with role tokens and response masking.
Processes each record (single-turn or multi-turn) and generates sequences
with role markers and masking (0=ignore user, 1=train on assistant).
Args:
records: List of JSONL record dicts with various SFT formats
Returns:
Dict with 'input_ids', 'attention_mask', and 'mask' lists, where each
element is a fixed-size sequence of length max_seq_length with masking info
"""
# Import here to avoid circular imports
from taoTrain.data.sft_utils import parse_sft_record, build_sft_sequence_tokens
max_seq_length = self.config.model.max_seq_length
user_token = getattr(self.config, 'user_token', '<user>')
assistant_token = getattr(self.config, 'assistant_token', '<assistant>')
sequences_input_ids = []
sequences_attention_masks = []
sequences_masks = []
for record in records:
try:
# Parse SFT record (supports multiple formats)
turns, is_multi_turn = parse_sft_record(record, self.config)
if not turns:
# Skip records that couldn't be parsed
continue
# Build token sequence with role tokens and response masking
input_ids, attention_mask, mask = build_sft_sequence_tokens(
turns=turns,
tokenizer=self.tokenizer,
user_token=user_token,
assistant_token=assistant_token,
max_seq_length=max_seq_length,
)
sequences_input_ids.append(input_ids)
sequences_attention_masks.append(attention_mask)
sequences_masks.append(mask)
except Exception as e:
# Log error but continue processing
print(f"Warning: Failed to tokenize SFT record: {e}")
continue
return {
"input_ids": sequences_input_ids,
"attention_mask": sequences_attention_masks,
"mask": sequences_masks,
}
def get_next_chunk(self, timeout: Optional[float] = None) -> Optional[Dict[str, List]]:
"""
Get the next tokenized chunk from the queue.
This is a blocking call that waits for the next chunk to be tokenized.
Returns None if queue is closed or all chunks have been processed.
CRITICAL: Always attempts to drain the queue first before returning None.
This prevents abandoning buffered chunks when threads finish.
Args:
timeout: Timeout in seconds (None = wait indefinitely)
Returns:
Dict with tokenized chunk, or None if queue is exhausted
Raises:
RuntimeError: If an error occurred in any worker thread
"""
if self._error_event.is_set():
error_summary = "; ".join(self._error_messages) if self._error_messages else "Unknown error"
raise RuntimeError(f"Tokenization thread error: {error_summary}")
# PRIORITY: Try to get from queue first (may have buffered items)
try:
chunk = self._queue.get(timeout=timeout)
return chunk
except queue.Empty:
# Queue is empty - check if threads are still working
with self._active_threads_lock:
if self._active_threads == 0 and self._next_chunk_idx >= len(self._chunk_order):
# All chunks assigned AND no active threads = true exhaustion
return None
# Queue temporarily empty but threads still working - signal to wait
return None
@property
def is_exhausted(self) -> bool:
"""Return True only when all chunks are assigned and all workers are idle."""
with self._active_threads_lock:
return self._active_threads == 0 and self._next_chunk_idx >= len(self._chunk_order)
def shutdown(self, wait: bool = True):
"""
Shutdown the tokenization worker threads gracefully.
Args:
wait: If True, wait for all threads to finish; otherwise return immediately
"""
if not self._threads:
return
# Signal threads to stop
self._stop_event.set()
# Drain queue to unblock threads if they're waiting to put
try:
while True:
self._queue.get_nowait()
except queue.Empty:
pass
# Wait for all threads to finish
if wait:
for thread in self._threads:
thread.join(timeout=5.0)
if thread.is_alive():
print(f"⚠ Tokenization thread {thread.name} did not terminate cleanly")
# Clear thread list to allow fresh start in next epoch
self._threads.clear()
print("✓ TokenizationQueue shutdown complete, thread list cleared")
def reset_for_next_epoch(self):
"""
Reset queue state for the next epoch.
This allows the same TokenizationQueue to be reused across multiple epochs.
Resets the chunk index counter, reshuffles chunks (if enabled), and clears
any buffered items and error state.
Called by AsyncBatchIterator at the start of epoch 2+.
"""
# Reset iteration counter
self._next_chunk_idx = 0
# Reshuffle chunk order if enabled
if self.shuffle_chunks:
import random
random.shuffle(self._chunk_order)
print(f"✓ Reshuffled chunk order for next epoch: {self._chunk_order}")
# Drain any remaining items from queue
items_drained = 0
try:
while True:
self._queue.get_nowait()
items_drained += 1
except queue.Empty:
pass
if items_drained > 0:
print(f"⚠ Drained {items_drained} items from queue before epoch reset")
# Clear error state
self._error_event.clear()
self._error_messages.clear()
# Clear threads list so new threads will be started in next epoch
self._threads.clear()
print(f"✓ TokenizationQueue reset for next epoch. Ready to process {len(self._chunk_order)} chunks")
def __len__(self) -> int:
"""Return total number of samples."""
return self.total_items
def __del__(self):
"""Cleanup on deletion."""
self.shutdown(wait=False)
|