import os # Set parallelism env var *before* importing tokenizers os.environ["TOKENIZERS_PARALLELISM"] = "false" import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader # Import necessary dataset functions, including concatenate_datasets if needed later from datasets import load_dataset, disable_caching, concatenate_datasets from tokenizers import Tokenizer, models, trainers, pre_tokenizers, processors, decoders import math import re from datetime import datetime from contextlib import nullcontext from collections import defaultdict import logging import random # For shuffling combined data # Disable caching for datasets if needed, helps ensure reprocessing # disable_caching() # Setup logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', force=True # Add this ) # Configuration CONFIG = { # --- Scaled Parameters --- "dim": 768, "n_layers": 16, "n_heads": 16, "ff_dim": 3072, # Explicitly set to 4 * dim # --- Kept Parameters --- "dropout": 0.1, "max_seq_len": 512, "vocab_size": 32000, # Fixed by tokenizer # --- Training/Dataset Parameters --- "batch_size": 12, "checkpoint_interval": 2000, "debug_interval": 400, # --- ADDED CoQA and QuAC --- "datasets": ["daily_dialog", "empathetic_dialogues", "blended_skill_talk", "AlekseyKorshuk/persona-chat"], "tokenizer_name": "hrom_tokenizer.json", # New name for expanded tokenizer "checkpoint_dir": "checkpoints", # Separate directory for expanded data model # --- Increased samples per dataset slightly for tokenizer --- "tokenizer_train_samples_per_dataset": 100000, # Use same limit for all, incl. new ones "learning_rate": 1e-5, "warmup_steps": 1000, "max_turns": 8, # Keep max_turns limit for Q&A datasets too "max_checkpoints": 5, "num_epochs": 30, "grad_accum_steps": 16 } # --- Model Definition (HROM, HROMBlock, HROMAttention, SwiGLU, RoPE) --- # (These classes remain unchanged from the previous version) class RotaryEmbedding(nn.Module): def __init__(self, dim): super().__init__() inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) def forward(self, seq_len): t = torch.arange(seq_len, device=self.inv_freq.device).type_as(self.inv_freq) freqs = torch.einsum("i, j -> i j", t, self.inv_freq) if seq_len == 0: return torch.empty((0, self.inv_freq.shape[0] * 2), device=self.inv_freq.device) # Defensive reshape only if necessary if freqs.shape[0] != seq_len and seq_len > 0: freqs = freqs.reshape(seq_len, -1) elif seq_len == 0: # Handle edge case for empty sequences return torch.empty((0, self.inv_freq.shape[0]*2), device=self.inv_freq.device, dtype=self.inv_freq.dtype) return torch.cat((freqs, freqs), dim=-1) def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(pos, t): # pos: (T, dim_rotary), t: (B, H, T, Head_Dim) pos = pos.to(t.device, dtype=t.dtype) pos = pos.unsqueeze(0).unsqueeze(1) # Shape: (1, 1, T, dim_rotary) tensor_seq_len = t.shape[2] pos_seq_len = pos.shape[2] if pos_seq_len < tensor_seq_len: logging.warning(f"RoPE Warning: pos sequence length ({pos_seq_len}) is shorter than tensor sequence length ({tensor_seq_len}). Using truncated tensor length for RoPE.") # This case is tricky, maybe only apply to the length of pos? # Or indicates an issue upstream. Let's slice t for now, though it's unusual. t_rotated = t[:, :, :pos_seq_len, :] pos = pos[:, :, :pos_seq_len, :] # Ensure pos matches the sliced tensor length # Apply rotation only to the slice cos_pos = pos.cos() sin_pos = pos.sin() t_rotated = (t_rotated * cos_pos) + (rotate_half(t_rotated) * sin_pos) # Concatenate the rotated part with the un-rotated part t_unrotated = t[:, :, pos_seq_len:, :] return torch.cat([t_rotated, t_unrotated], dim=2) elif pos_seq_len > tensor_seq_len: pos = pos[:, :, :tensor_seq_len, :] # Slice pos to match tensor # Check dimension match after potential slicing if pos.shape[-1] != t.shape[-1]: logging.error(f"Mismatched dimensions for RoPE: pos ({pos.shape[-1]}) vs t ({t.shape[-1]})") raise ValueError("Rotary embedding dimension must match head dimension.") cos_pos = pos.cos() sin_pos = pos.sin() rotated_t = (t * cos_pos) + (rotate_half(t) * sin_pos) return rotated_t class SwiGLU(nn.Module): def forward(self, x): x, gate = x.chunk(2, dim=-1) return x * nn.functional.gelu(gate) class HROMAttention(nn.Module): def __init__(self): super().__init__() self.dim = CONFIG["dim"] self.n_heads = CONFIG["n_heads"] self.head_dim = self.dim // self.n_heads if self.dim % self.n_heads != 0: raise ValueError("dim must be divisible by n_heads") self.qkv = nn.Linear(self.dim, 3 * self.dim) self.proj = nn.Linear(self.dim, self.dim) self.rotary = RotaryEmbedding(self.head_dim) self.dropout = nn.Dropout(CONFIG["dropout"]) def forward(self, x, mask=None): B, T, C = x.shape qkv = self.qkv(x) qkv = qkv.reshape(B, T, 3, self.n_heads, self.head_dim) q, k, v = qkv.unbind(2) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) # Generate RoPE embeddings for the current sequence length T pos = self.rotary(T) # Shape (T, Head_Dim) # Apply RoPE q = apply_rotary_pos_emb(pos, q) k = apply_rotary_pos_emb(pos, k) # Attention calculation attn_scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim)) if mask is not None: # Ensure mask is broadcastable (B, 1, T, T) if mask.dim() == 2: # (B, T) -> (B, 1, 1, T) -> add with causal = (B, 1, T, T) mask = mask.unsqueeze(1).unsqueeze(2) elif mask.dim() == 3: # (B, T, T) mask = mask.unsqueeze(1) # Add mask AFTER scaling scores attn_scores = attn_scores + mask # Add large negative values for masked positions # Softmax and dropout attn_probs = torch.softmax(attn_scores.float(), dim=-1).to(dtype=x.dtype) # Use float for stability attn_probs = self.dropout(attn_probs) # Output projection output = attn_probs @ v output = output.transpose(1, 2).reshape(B, T, self.dim) return self.proj(output) class HROMBlock(nn.Module): def __init__(self): super().__init__() self.attn = HROMAttention() self.ff = nn.Sequential( nn.Linear(CONFIG["dim"], 2 * CONFIG["ff_dim"]), SwiGLU(), nn.Linear(CONFIG["ff_dim"], CONFIG["dim"]) ) self.norm1 = nn.LayerNorm(CONFIG["dim"]) self.norm2 = nn.LayerNorm(CONFIG["dim"]) self.dropout = nn.Dropout(CONFIG["dropout"]) def forward(self, x, mask=None): # Pre-Normalization normed_x = self.norm1(x) attn_output = self.attn(normed_x, mask) x = x + self.dropout(attn_output) normed_x = self.norm2(x) ff_output = self.ff(normed_x) x = x + self.dropout(ff_output) return x class HROM(nn.Module): def __init__(self): super().__init__() self.embed = nn.Embedding(CONFIG["vocab_size"], CONFIG["dim"]) self.blocks = nn.ModuleList([HROMBlock() for _ in range(CONFIG["n_layers"])]) self.norm = nn.LayerNorm(CONFIG["dim"]) self.head = nn.Linear(CONFIG["dim"], CONFIG["vocab_size"]) self.dropout = nn.Dropout(CONFIG["dropout"]) # Add dropout after embedding self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.xavier_uniform_(module.weight) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.LayerNorm): torch.nn.init.zeros_(module.bias) torch.nn.init.ones_(module.weight) def forward(self, input_ids, attention_mask=None): B, T = input_ids.shape x = self.embed(input_ids) x = self.dropout(x) # Apply dropout after embedding # Create the combined mask for attention combined_mask = None # Start with causal mask valid for all sequences in batch causal_mask = torch.triu(torch.ones(T, T, device=input_ids.device) * float('-inf'), diagonal=1) combined_mask = causal_mask.unsqueeze(0).unsqueeze(1) # (1, 1, T, T) if attention_mask is not None: # Process padding mask from attention_mask (0 = pad, 1 = real) # Convert 0s to -inf, 1s to 0 pad_mask = (1.0 - attention_mask.to(torch.float32)) * torch.finfo(torch.float32).min pad_mask = pad_mask.unsqueeze(1).unsqueeze(2) # (B, 1, 1, T) # Add padding mask to causal mask. Broadcasting ensures (B, 1, T, T) # Where pad_mask is -inf, the result is -inf. Otherwise, it's the causal value. combined_mask = combined_mask + pad_mask # Ensure mask dtype matches data dtype (esp. for AMP) combined_mask = combined_mask.to(dtype=x.dtype) for block in self.blocks: x = block(x, combined_mask) # Pass the combined mask to each block x = self.norm(x) logits = self.head(x) return logits # --- Tokenizer Training --- class TokenizerTrainer: def __init__(self): self.tokenizer = Tokenizer(models.BPE()) self.tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) self.tokenizer.decoder = decoders.ByteLevel() self.special_tokens = ["", "", "", "", "", ""] # Use the updated tokenizer name from CONFIG self.tokenizer_path = os.path.join("tokenizer", CONFIG["tokenizer_name"]) self.tokenizer_dir = os.path.dirname(self.tokenizer_path) def _clean_text(self, text): text = str(text) # Ensure text is string text = re.sub(r'_comma_', ',', text) # Allow alphanumeric, whitespace, and basic punctuation including quotes text = re.sub(r'[^\w\s.,!?\'\-:;<>"]', '', text) text = re.sub(r'\s+', ' ', text).strip() return text def train(self, dataset_names): logging.info("Starting tokenizer training...") text_samples = [] samples_per_dataset = CONFIG['tokenizer_train_samples_per_dataset'] # --- Process DailyDialog --- if "daily_dialog" in dataset_names: logging.info(f"Loading daily_dialog for tokenizer training (max {samples_per_dataset} dialogues)...") try: # Limit dialogues loaded directly using slicing dd_dataset = load_dataset("daily_dialog", split=f"train[:{samples_per_dataset}]", trust_remote_code=True) # Add trust_remote_code=True logging.info("Processing daily_dialog...") for entry in dd_dataset: formatted_dialogue = [] dialogue = entry['dialog'][:CONFIG["max_turns"]] for i, utterance in enumerate(dialogue): role = "" if i % 2 == 0 else "" cleaned_utterance = self._clean_text(utterance) if cleaned_utterance: # Only add non-empty turns formatted_dialogue.append(f"{role} {cleaned_utterance}") if formatted_dialogue: # Only add if dialogue is not empty after cleaning text_samples.append(" ".join(formatted_dialogue)) except Exception as e: logging.error(f"Failed to load or process daily_dialog for tokenizer: {e}") # --- Process EmpatheticDialogues --- if "empathetic_dialogues" in dataset_names: logging.info(f"Loading empathetic_dialogues for tokenizer training (max {samples_per_dataset} dialogues)...") try: # Load more initially to ensure we get enough unique conversations (adjust multiplier if needed) ed_dataset = load_dataset("empathetic_dialogues", split=f"train[:{samples_per_dataset * 3}]", trust_remote_code=True) # Add trust_remote_code=True logging.info("Processing empathetic_dialogues...") conversations = defaultdict(list) processed_conv_count = 0 # Group utterances by conv_id first grouped_by_conv = defaultdict(list) for entry in ed_dataset: grouped_by_conv[entry['conv_id']].append(entry) # Process conversations ensuring max samples limit for conv_id, entries in grouped_by_conv.items(): if processed_conv_count >= samples_per_dataset: break # Sort by utterance_idx to maintain order sorted_entries = sorted(entries, key=lambda x: x['utterance_idx']) formatted_dialogue = [] # Handle context and first utterance if sorted_entries[0]['context']: cleaned_context = self._clean_text(sorted_entries[0]['context']) if cleaned_context: formatted_dialogue.append(f" {cleaned_context}") # Assume context is user start # Process subsequent utterances last_role = '' if formatted_dialogue else None # Set initial last role based on context for entry in sorted_entries: cleaned_utterance = self._clean_text(entry['utterance']) if cleaned_utterance: # Determine role based on alternation current_role = '' if last_role == '' else '' formatted_dialogue.append(f"{current_role} {cleaned_utterance}") last_role = current_role # Update last role # Apply max turns limit to the formatted turns formatted_dialogue = formatted_dialogue[:CONFIG["max_turns"]] if formatted_dialogue: text_samples.append(" ".join(formatted_dialogue)) processed_conv_count += 1 # Count processed unique conversations except Exception as e: logging.error(f"Failed to load or process empathetic_dialogues for tokenizer: {e}") # --- Process BlendedSkillTalk --- if "blended_skill_talk" in dataset_names: logging.info(f"Loading blended_skill_talk for tokenizer training (max {samples_per_dataset} dialogues)...") try: # Load dialogues - BST is structured differently, slice directly bst_dataset = load_dataset("blended_skill_talk", split=f"train[:{samples_per_dataset}]", trust_remote_code=True) # Add trust_remote_code=True logging.info("Processing blended_skill_talk...") for entry in bst_dataset: formatted_dialogue = [] # Combine the dialogue history and the final two turns dialogue_turns_raw = entry['previous_utterance'] # Add final utterances if they exist and are not empty strings if entry.get('free_turker_utterance'): dialogue_turns_raw.append(entry['free_turker_utterance']) if entry.get('guided_turker_utterance'): dialogue_turns_raw.append(entry['guided_turker_utterance']) turns_to_process = dialogue_turns_raw[:CONFIG["max_turns"]] # Apply max turns limit for i, utterance in enumerate(turns_to_process): role = "" if i % 2 == 0 else "" # Assume simple alternation cleaned_utterance = self._clean_text(utterance) if cleaned_utterance: formatted_dialogue.append(f"{role} {cleaned_utterance}") if formatted_dialogue: text_samples.append(" ".join(formatted_dialogue)) except Exception as e: logging.error(f"Failed to load or process blended_skill_talk for tokenizer: {e}") # --- Process PersonaChat --- if "AlekseyKorshuk/persona-chat" in dataset_names: # Correct dataset identifier pc_dataset_name = "AlekseyKorshuk/persona-chat" logging.info(f"Loading {pc_dataset_name} for tokenizer training (max {samples_per_dataset} dialogues)...") try: pc_dataset = load_dataset(pc_dataset_name, split=f"train[:{samples_per_dataset}]", trust_remote_code=True) # Add trust_remote_code=True, Correct dataset identifier logging.info(f"Processing {pc_dataset_name}...") for entry in pc_dataset: # PersonaChat often has 'utterances' containing 'history' if 'utterances' in entry and entry['utterances']: # Get the history from the last item in utterances for the full dialogue history = entry['utterances'][-1]['history'] history = history[:CONFIG["max_turns"]] # Apply max turns formatted_dialogue = [] for i, utterance in enumerate(history): role = "" if i % 2 == 0 else "" # Assume simple alternation cleaned_utterance = self._clean_text(utterance) if cleaned_utterance: formatted_dialogue.append(f"{role} {cleaned_utterance}") if formatted_dialogue: text_samples.append(" ".join(formatted_dialogue)) else: logging.warning(f"Skipping {pc_dataset_name} entry due to unexpected structure: {entry}") except Exception as e: logging.error(f"Failed to load or process {pc_dataset_name} for tokenizer: {e}") logging.info(f"Total text samples for tokenizer training: {len(text_samples)}") if not text_samples: raise ValueError("No text samples collected for tokenizer training. Check dataset loading and paths.") # Ensure tokenizer directory exists before training os.makedirs(self.tokenizer_dir, exist_ok=True) logging.info(f"Training BPE tokenizer with vocab size {CONFIG['vocab_size']}...") trainer = trainers.BpeTrainer( vocab_size=CONFIG["vocab_size"], special_tokens=self.special_tokens, min_frequency=2, # Keep min_frequency low with more data show_progress=True ) # Make sure text_samples is an iterator or list of strings def text_iterator(): for sample in text_samples: yield sample self.tokenizer.train_from_iterator(text_iterator(), trainer=trainer, length=len(text_samples)) eos_token_id = self.tokenizer.token_to_id("") if eos_token_id is None: logging.warning(" token not found in trained tokenizer vocab! Using as fallback for post-processor.") eos_token_id = self.tokenizer.token_to_id("") or 0 # Fallback needed # Configure post-processor (adjust if needed based on how you structure input/output) self.tokenizer.post_processor = processors.TemplateProcessing( single="$A ", pair="$A $B ", # How to handle pairs - maybe just use single always? special_tokens=[("", eos_token_id)], ) logging.info(f"Saving tokenizer to {self.tokenizer_path}") self.tokenizer.save(self.tokenizer_path) logging.info("Tokenizer training complete.") def get_tokenizer(self): if not os.path.exists(self.tokenizer_path): raise FileNotFoundError(f"Tokenizer file not found at {self.tokenizer_path}. Train tokenizer first.") tokenizer = Tokenizer.from_file(self.tokenizer_path) # Verify special tokens crucial for processing exist required_tokens = ["", "", "", "", "", ""] for token in required_tokens: if tokenizer.token_to_id(token) is None: raise ValueError(f"Crucial special token '{token}' not found in loaded tokenizer '{self.tokenizer_path}'!") return tokenizer # --- Dataset Loading and Processing --- class CombinedChatDataset(Dataset): def __init__(self, tokenizer): self.tokenizer = tokenizer self.pad_id = self.tokenizer.token_to_id("") self.eos_id = self.tokenizer.token_to_id("") self.bos_id = self.tokenizer.token_to_id("") self.user_id = self.tokenizer.token_to_id("") self.assistant_id = self.tokenizer.token_to_id("") self.max_length = CONFIG["max_seq_len"] # Reuse cleaning function from TokenizerTrainer instance self._clean_text = TokenizerTrainer()._clean_text self.all_processed_conversations = [] # --- Process DailyDialog --- if "daily_dialog" in CONFIG["datasets"]: logging.info("Loading and processing daily_dialog dataset...") try: dd_dataset = load_dataset("daily_dialog", split="train", trust_remote_code=True) # Add trust_remote_code=True logging.info(f"Processing {len(dd_dataset)} daily_dialog conversations...") for entry in dd_dataset: conversation = [] dialogue = entry['dialog'][:CONFIG["max_turns"]] if not dialogue: continue for i, utterance in enumerate(dialogue): role = "" if i % 2 == 0 else "" cleaned_text = self._clean_text(utterance) if cleaned_text: conversation.append({'role': role, 'text': cleaned_text}) if conversation: self.all_processed_conversations.append(conversation) except Exception as e: logging.error(f"Failed to load or process daily_dialog for training: {e}") # --- Process EmpatheticDialogues --- if "empathetic_dialogues" in CONFIG["datasets"]: logging.info("Loading and processing empathetic_dialogues dataset...") try: ed_dataset = load_dataset("empathetic_dialogues", split="train", trust_remote_code=True) # Add trust_remote_code=True logging.info("Grouping empathetic_dialogues by conversation ID...") conversations_grouped = defaultdict(list) for entry in ed_dataset: conversations_grouped[entry['conv_id']].append(entry) logging.info(f"Processing {len(conversations_grouped)} empathetic_dialogues conversations...") for conv_id, entries in conversations_grouped.items(): conversation = [] sorted_entries = sorted(entries, key=lambda x: x['utterance_idx']) # Handle context as first user turn if present if sorted_entries[0]['context']: context_text = self._clean_text(sorted_entries[0]['context']) if context_text: conversation.append({'role': '', 'text': context_text}) # Process utterances, assuming alternation last_role = conversation[-1]['role'] if conversation else None # Role of the last added turn for entry in sorted_entries: text = self._clean_text(entry['utterance']) if not text: continue # Determine role based on the *last added* role current_role = '' if last_role == '' else '' conversation.append({'role': current_role, 'text': text}) last_role = current_role # Update for next iteration # Apply max turns limit *after* forming the full sequence conversation = conversation[:CONFIG["max_turns"]] if conversation: self.all_processed_conversations.append(conversation) except Exception as e: logging.error(f"Failed to load or process empathetic_dialogues for training: {e}") # --- Process BlendedSkillTalk --- if "blended_skill_talk" in CONFIG["datasets"]: logging.info("Loading and processing blended_skill_talk dataset...") try: bst_dataset = load_dataset("blended_skill_talk", split="train", trust_remote_code=True) # Add trust_remote_code=True logging.info(f"Processing {len(bst_dataset)} blended_skill_talk conversations...") for entry in bst_dataset: conversation = [] # Reconstruct dialogue: history + final two turns (if they exist) dialogue_turns_raw = entry['previous_utterance'] if entry.get('free_turker_utterance'): dialogue_turns_raw.append(entry['free_turker_utterance']) if entry.get('guided_turker_utterance'): dialogue_turns_raw.append(entry['guided_turker_utterance']) if not dialogue_turns_raw: continue # Skip if no turns found turns_to_process = dialogue_turns_raw[:CONFIG["max_turns"]] # Apply max turns limit for i, utterance in enumerate(turns_to_process): role = "" if i % 2 == 0 else "" # Assume simple alternation cleaned_text = self._clean_text(utterance) if cleaned_text: conversation.append({'role': role, 'text': cleaned_text}) if conversation: # Only add if not empty after cleaning/truncation self.all_processed_conversations.append(conversation) except Exception as e: logging.error(f"Failed to load or process blended_skill_talk for training: {e}") # --- Process PersonaChat --- if "AlekseyKorshuk/persona-chat" in CONFIG["datasets"]: # Correct dataset identifier pc_dataset_name = "AlekseyKorshuk/persona-chat" logging.info(f"Loading and processing {pc_dataset_name} dataset...") try: pc_dataset = load_dataset(pc_dataset_name, split="train", trust_remote_code=True) # Add trust_remote_code=True, Correct dataset identifier logging.info(f"Processing {len(pc_dataset)} {pc_dataset_name} conversations...") for entry in pc_dataset: conversation = [] if 'utterances' in entry and entry['utterances']: # Extract the dialogue history history = entry['utterances'][-1]['history'] history = history[:CONFIG["max_turns"]] # Apply max turns limit for i, utterance in enumerate(history): role = "" if i % 2 == 0 else "" # Simple alternation cleaned_text = self._clean_text(utterance) if cleaned_text: conversation.append({'role': role, 'text': cleaned_text}) if conversation: # Only add if not empty self.all_processed_conversations.append(conversation) else: logging.warning(f"Skipping {pc_dataset_name} entry due to unexpected structure: {entry.keys()}") except Exception as e: logging.error(f"Failed to load or process {pc_dataset_name} for training: {e}") logging.info(f"Total processed conversations from all datasets: {len(self.all_processed_conversations)}") if not self.all_processed_conversations: raise ValueError("No processed conversations were created from any dataset. Check loading logic and dataset availability.") logging.info("Shuffling combined dataset...") random.shuffle(self.all_processed_conversations) def __len__(self): return len(self.all_processed_conversations) def __getitem__(self, idx): conversation = self.all_processed_conversations[idx] formatted_ids = [self.bos_id] for turn in conversation: role_id = self.user_id if turn['role'] == '' else self.assistant_id # Encode without adding special tokens automatically by tokenizer try: utterance_ids = self.tokenizer.encode(turn['text'], add_special_tokens=False).ids except Exception as e: logging.error(f"Error encoding text at index {idx}, turn '{turn}': {e}") utterance_ids = [] # Skip this utterance on error # Check length: Current + Role + Utterance + EOS <= MaxLength # Need +1 for role, +len(utterance), +1 for potential EOS if len(formatted_ids) + 1 + len(utterance_ids) + 1 > self.max_length: # Attempt to add just the role and EOS if utterance is too long if len(formatted_ids) + 1 + 1 <= self.max_length: formatted_ids.append(role_id) formatted_ids.append(self.eos_id) break # Stop adding turns formatted_ids.append(role_id) formatted_ids.extend(utterance_ids) formatted_ids.append(self.eos_id) # Final safety truncate (should be rare if logic above is correct) if len(formatted_ids) > self.max_length: formatted_ids = formatted_ids[:self.max_length] # Ensure last token isn't partial (though unlikely with BPE) # If the truncated sequence ends with a role ID, it's probably bad, remove it. if formatted_ids and (formatted_ids[-1] == self.user_id or formatted_ids[-1] == self.assistant_id): formatted_ids.pop() # If after popping the role ID, it's still too long (unlikely), truncate again if len(formatted_ids) > self.max_length: formatted_ids = formatted_ids[:self.max_length] # Handle case of extremely short sequences after processing if len(formatted_ids) < 2: # Need at least BOS and one other token for input/label pair logging.warning(f"Sequence at index {idx} is too short after processing (<2 tokens). Skipping. Original length: {len(conversation)}") # Return None to be filtered by collate_fn return None input_ids = formatted_ids[:-1] labels = formatted_ids[1:] # Final check before returning if len(input_ids) == 0: logging.warning(f"Sequence at index {idx} resulted in empty input_ids after slicing. Skipping.") return None return {"input_ids": input_ids, "labels": labels} @staticmethod def collate_fn(batch): # Filter out None items from __getitem__ batch = [item for item in batch if item is not None] if not batch: return None # Return None if the whole batch was invalid max_len = max(len(item["input_ids"]) for item in batch) # Load tokenizer once to get pad_id - ensure path matches CONFIG try: # Correctly reference the tokenizer path from CONFIG within the static method tokenizer_path = os.path.join("tokenizer", CONFIG["tokenizer_name"]) # TODO: Consider passing tokenizer/pad_id if this becomes a bottleneck tokenizer = Tokenizer.from_file(tokenizer_path) pad_id = tokenizer.token_to_id("") if pad_id is None: raise ValueError(" token not found") except Exception as e: logging.error(f"Collate Error: Failed to load tokenizer or get pad_id ('{CONFIG['tokenizer_name']}'): {e}") pad_id = 0 # Risky fallback inputs, labels, masks = [], [], [] for item in batch: input_len = len(item["input_ids"]) pad_len = max_len - input_len inputs.append(item["input_ids"] + [pad_id] * pad_len) # Pad labels with pad_id (or any ID to be ignored by CrossEntropyLoss) labels.append(item["labels"] + [pad_id] * pad_len) masks.append([1] * input_len + [0] * pad_len) return { "input_ids": torch.tensor(inputs, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), "attention_mask": torch.tensor(masks, dtype=torch.long) # Or bool } # --- Trainer, Safety Manager, Checkpoint Manager --- class HROMTrainer: def __init__(self, model, tokenizer): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logging.info(f"Using device: {self.device}") self.model = model.to(self.device) self.use_amp = (self.device.type == "cuda" and hasattr(torch.cuda.amp, "GradScaler")) self.scaler = torch.cuda.amp.GradScaler() if self.use_amp else None logging.info(f"Automatic Mixed Precision (AMP): {'Enabled' if self.use_amp else 'Disabled'}") self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=CONFIG["learning_rate"], # Base LR betas=(0.9, 0.95), weight_decay=0.1, fused= (self.device.type == "cuda") ) self.tokenizer = tokenizer self.pad_id = self.tokenizer.token_to_id("") if self.pad_id is None: # Attempt to get from config if available or fallback self.pad_id = CONFIG.get("pad_token_id", 0) logging.warning(f" token ID not found in tokenizer, using fallback ID: {self.pad_id}") # Make sure ignore_index uses the determined pad_id self.criterion = nn.CrossEntropyLoss(ignore_index=self.pad_id) self.base_lr = CONFIG["learning_rate"] self.warmup_steps = CONFIG["warmup_steps"] def _adjust_learning_rate(self, step): if self.warmup_steps > 0 and step < self.warmup_steps: lr = self.base_lr * (step + 1) / self.warmup_steps else: # Optional: Add LR decay (e.g., cosine) after warmup # Example: lr = self.base_lr * 0.5 * (1 + math.cos(math.pi * (step - self.warmup_steps) / (total_steps - self.warmup_steps))) lr = self.base_lr # Keep base LR after warmup for now for param_group in self.optimizer.param_groups: param_group['lr'] = lr return lr def train_step(self, batch): # Determine precision for autocast if self.use_amp: amp_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16 autocast_context = torch.cuda.amp.autocast(dtype=amp_dtype, enabled=self.use_amp) if self.use_amp else nullcontext() with autocast_context: input_ids = batch["input_ids"].to(self.device) attention_mask = batch["attention_mask"].to(self.device) labels = batch["labels"].to(self.device) outputs = self.model(input_ids, attention_mask=attention_mask) # Reshape for loss calculation logits_flat = outputs.view(-1, outputs.size(-1)) # Shape: (B * T, vocab_size) labels_flat = labels.view(-1) # Shape: (B * T) # Calculate loss - ensure logits are float32 for stability esp. with AMP loss = self.criterion(logits_flat.float(), labels_flat) # Scale loss for gradient accumulation scaled_loss = loss / CONFIG["grad_accum_steps"] # Backward pass if self.use_amp and self.scaler: self.scaler.scale(scaled_loss).backward() else: scaled_loss.backward() return loss.item() # Return the unscaled loss for logging def clip_and_step(self, current_optimizer_step): current_lr = self._adjust_learning_rate(current_optimizer_step) # Gradient Clipping *before* optimizer step if self.use_amp and self.scaler: # Unscale first - important before clipping self.scaler.unscale_(self.optimizer) # Clip grad norm torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) # Optimizer step (with scaler) self.scaler.step(self.optimizer) # Update scaler for next iteration self.scaler.update() else: # Clip grad norm torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0) # Optimizer step self.optimizer.step() # Zero gradients *after* stepping self.optimizer.zero_grad(set_to_none=True) return current_lr class SafetyManager: # (No changes needed in SafetyManager implementation itself) def __init__(self, model, tokenizer): self.model = model self.tokenizer = tokenizer # More conservative list self.bad_words = ["kill", "murder", "suicide", "hate", "abuse", "violence", "illegal", "harm", "die", "attack", "rape", "molest", "exploit", "terror"] self.bad_word_ids = [] logging.info("Initializing safety manager...") # Pre-encode bad word sequences for word in self.bad_words: # Encode potentially multi-token words carefully ids = tokenizer.encode(f" {word}", add_special_tokens=False).ids # Add prefix space for BPE if ids: self.bad_word_ids.append(ids) logging.debug(f"Encoded bad word '{word}' (with space) to IDs: {ids}") # Try without space too ids_no_space = tokenizer.encode(word, add_special_tokens=False).ids if ids_no_space and ids_no_space != ids: self.bad_word_ids.append(ids_no_space) logging.debug(f"Encoded bad word '{word}' (no space) to IDs: {ids_no_space}") if not ids and not ids_no_space: logging.warning(f"Could not encode bad word '{word}' - skipping.") # Pre-get special IDs self.eos_id = self.tokenizer.token_to_id("") self.bos_id = self.tokenizer.token_to_id("") self.user_id = self.tokenizer.token_to_id("") self.assistant_id = self.tokenizer.token_to_id("") self.pad_id = self.tokenizer.token_to_id("") if self.eos_id is None: logging.error(" token ID not found for SafetyManager!"); self.eos_id = 0 if self.bos_id is None: logging.error(" token ID not found for SafetyManager!"); self.bos_id = 0 if self.user_id is None: logging.error(" token ID not found for SafetyManager!") if self.assistant_id is None: logging.error(" token ID not found for SafetyManager!") if self.pad_id is None: logging.error(" token ID not found for SafetyManager!"); self.pad_id = 0 def contains_sequence(self, tokens, seq): """Checks if the list `tokens` contains the sublist `seq`.""" if not seq or not tokens or len(tokens) < len(seq): return False seq_len = len(seq) for i in range(len(tokens) - seq_len + 1): if tokens[i : i + seq_len] == seq: return True return False def content_filter(self, text_ids): """Checks if a list of token IDs contains any bad word sequences.""" if not isinstance(text_ids, list): logging.warning("Content filter received non-list input.") return True # Default to safe if input is weird for bad_ids in self.bad_word_ids: if self.contains_sequence(text_ids, bad_ids): # Log the detected sequence for debugging detected_word = self.tokenizer.decode(bad_ids) logging.warning(f"Unsafe content detected: Found sequence corresponding to '{detected_word}' (IDs: {bad_ids}).") return False # Unsafe return True # Safe def generate_safely(self, prompt, max_new_tokens=50, temperature=0.5, top_k=50): self.model.eval() device = next(self.model.parameters()).device # Encode prompt, ensure it ends appropriately (e.g., with role token + EOS?) # Let's assume the prompt ends like " blah blah " and we need to add "" prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids # Start generation sequence with BOS, prompt, and assistant token # Ensure prompt doesn't already include BOS if prompt_ids and prompt_ids[0] == self.bos_id: input_ids = list(prompt_ids) else: input_ids = [self.bos_id] + list(prompt_ids) # Add the assistant token to signal the model to generate the response if self.assistant_id is not None: input_ids.append(self.assistant_id) else: logging.error("Assistant token ID is None, cannot properly start generation.") return "Error: Assistant token not found." generated_ids = list(input_ids) # Start with the prepared input sequence logging.debug(f"Starting safe generation with initial IDs: {generated_ids}") with torch.no_grad(): for step in range(max_new_tokens): # Prepare input tensor for this step - only use up to max_seq_len current_input_ids = generated_ids[-CONFIG["max_seq_len"]:] current_input_tensor = torch.tensor([current_input_ids]).to(device) # Create attention mask for the current length attention_mask = torch.ones_like(current_input_tensor) # Model forward pass try: outputs = self.model(current_input_tensor, attention_mask=attention_mask) next_token_logits = outputs[:, -1, :] # Logits for the next token except Exception as e: logging.error(f"Model forward pass failed during generation: {e}") break # Stop generation on error # --- Safety Check BEFORE sampling --- # Apply penalties to bad word starting tokens if possible # For now, we filter *after* sampling the token # Sampling (Temperature, Top-K) if temperature > 0 and temperature != 1.0: next_token_logits = next_token_logits / temperature if top_k > 0 and top_k < next_token_logits.size(-1): # Ensure top_k is valid v, _ = torch.topk(next_token_logits, top_k) # Handle potential NaN/Inf in logits before comparison safe_logits = torch.nan_to_num(next_token_logits, nan=-float('inf'), posinf=float('inf'), neginf=-float('inf')) threshold = v[:, [-1]] safe_logits[safe_logits < threshold] = -float('Inf') next_token_logits = safe_logits # Use the filtered logits probs = torch.softmax(next_token_logits, dim=-1) # Handle potential NaNs in probabilities before sampling if torch.isnan(probs).any(): logging.warning("NaN detected in probabilities before sampling. Replacing with uniform distribution.") probs = torch.ones_like(probs) / probs.size(-1) # Fallback to uniform next_token_id = torch.multinomial(probs, num_samples=1).item() # --- Safety Check AFTER sampling token --- # Check if adding this token creates a bad sequence potential_sequence_ids = generated_ids + [next_token_id] # Check only the newly formed part for bad words for efficiency? # Let's check the whole sequence for simplicity/robustness for now. if not self.content_filter(potential_sequence_ids): logging.warning(f"Potential unsafe token ({next_token_id}, '{self.tokenizer.decode([next_token_id])}') blocked POST-sampling. Stopping generation.") # Optionally try sampling a different token? For now, just stop. break # Add the safe token generated_ids.append(next_token_id) # Check for EOS token if next_token_id == self.eos_id: logging.debug(f"EOS token generated at step {step+1}. Stopping generation.") break # Prevent infinite loops if max tokens reached if step == max_new_tokens - 1: logging.debug("Max new tokens reached. Stopping generation.") # Ensure the sequence ends with EOS if it didn't naturally if generated_ids[-1] != self.eos_id and self.eos_id is not None: generated_ids.append(self.eos_id) self.model.train() # Set model back to training mode # Decode the generated part (excluding the initial prompt + assistant token) start_index = len(input_ids) response_ids = generated_ids[start_index:] # Decode, skipping special tokens like EOS, BOS, PAD but potentially keeping USER/ASSISTANT # Let's skip all special tokens for the final output text for clarity. decoded_text = self.tokenizer.decode(response_ids, skip_special_tokens=True).strip() return decoded_text def debug_generation(self, prompt=" Tell me about your hobbies."): # Example prompt logging.info(f"\n--- Debug Generation & Safety Check ---") # Ensure prompt ends logically for the model (e.g., with user token and EOS) if not prompt.strip().endswith(""): if not prompt.strip().endswith("") and not prompt.strip().endswith(""): prompt = prompt.strip() + " " # Add EOS if ends mid-sentence else: prompt = prompt.strip() + " " # Add EOS after role token # Ensure the prompt starts appropriately (e.g., no BOS needed here as generate_safely adds it) if prompt.startswith(""): prompt = prompt[len(""):].strip() generated_response = self.generate_safely(prompt, max_new_tokens=60, temperature=0.7, top_k=50) logging.info(f"Prompt Sent: '{prompt}'") logging.info(f"Generated Response: '{generated_response}'") logging.info("\n--- End Debug Generation ---\n") class CheckpointManager: def __init__(self): # Use checkpoint directory from CONFIG self.checkpoint_dir = CONFIG["checkpoint_dir"] os.makedirs(self.checkpoint_dir, exist_ok=True) logging.info(f"Checkpoint directory set to: {self.checkpoint_dir}") def save(self, model, optimizer, step): timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") # Use a consistent naming scheme based on the directory name if desired prefix = os.path.basename(self.checkpoint_dir).replace("checkpoints_", "") # Ensure step is converted to string if it's passed as something else (e.g., 'final') step_str = str(step) filename = f"hrom_{prefix}_step{step_str}_{timestamp}.pt" path = os.path.join(self.checkpoint_dir, filename) state = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "step": step if isinstance(step, int) else -1, # Store step number or -1 for non-numeric steps "config": CONFIG # Save config with checkpoint } logging.info(f"Saving checkpoint to {path}...") try: torch.save(state, path) logging.info(f"Checkpoint saved successfully at step {step_str}.") self._cleanup_old_checkpoints() except Exception as e: logging.error(f"Failed to save checkpoint '{path}': {e}") def _cleanup_old_checkpoints(self): max_checkpoints = CONFIG.get("max_checkpoints", 5) # Get from config, default 5 if max_checkpoints <= 0: return # Keep all checkpoints if max_checkpoints is non-positive try: # Filter only files matching the expected pattern (avoid deleting other files) prefix = os.path.basename(self.checkpoint_dir).replace("checkpoints_", "") pattern = re.compile(rf"hrom_{prefix}_step(\d+|.+)_(\d{{8}}_\d{{6}})\.pt") checkpoints = [] for f in os.listdir(self.checkpoint_dir): match = pattern.match(f) if match: filepath = os.path.join(self.checkpoint_dir, f) checkpoints.append((filepath, os.path.getmtime(filepath))) # Sort by modification time (oldest first) checkpoints.sort(key=lambda x: x[1]) num_to_delete = len(checkpoints) - max_checkpoints if num_to_delete > 0: #logging.info(f"Max checkpoints ({max_checkpoints}) reached. Removing {num_to_delete} oldest checkpoints.") for i in range(num_to_delete): file_to_remove, _ = checkpoints[i] try: os.remove(file_to_remove) #logging.info(f"Removed old checkpoint: {os.path.basename(file_to_remove)}") except OSError as e: logging.error(f"Error removing checkpoint {file_to_remove}: {e}") except Exception as e: logging.error(f"Error during checkpoint cleanup: {e}") def load_latest(self, model, optimizer): try: # Filter files based on pattern and sort by time prefix = os.path.basename(self.checkpoint_dir).replace("checkpoints_", "") pattern = re.compile(rf"hrom_{prefix}_step(\d+|.+)_(\d{{8}}_\d{{6}})\.pt") checkpoints = [] for f in os.listdir(self.checkpoint_dir): match = pattern.match(f) if match: filepath = os.path.join(self.checkpoint_dir, f) checkpoints.append((filepath, os.path.getmtime(filepath))) if not checkpoints: logging.info("No valid checkpoints found to load.") return 0 # Start from step 0 # Sort by modification time (newest first) checkpoints.sort(key=lambda x: x[1], reverse=True) latest_checkpoint_path, _ = checkpoints[0] logging.info(f"Loading latest checkpoint from: {latest_checkpoint_path}") map_location = torch.device('cuda' if torch.cuda.is_available() else 'cpu') checkpoint = torch.load(latest_checkpoint_path, map_location=map_location) # --- Config Compatibility Check (Optional but Recommended) --- loaded_config = checkpoint.get("config", {}) # Compare key parameters that affect model architecture or data processing critical_keys = ["dim", "n_layers", "n_heads", "ff_dim", "vocab_size", "max_seq_len", "tokenizer_name"] mismatched_keys = [] if loaded_config: for key in critical_keys: # Check if key exists in both and if they differ if key in loaded_config and key in CONFIG and loaded_config[key] != CONFIG[key]: mismatched_keys.append((key, loaded_config[key], CONFIG[key])) # Check if key missing in current config but present in checkpoint elif key in loaded_config and key not in CONFIG: mismatched_keys.append((key, loaded_config[key], "Not in current CONFIG")) # Check if key missing in checkpoint config but present in current elif key not in loaded_config and key in CONFIG: mismatched_keys.append((key, "Not in loaded CONFIG", CONFIG[key])) if mismatched_keys: logging.warning("--- CONFIG MISMATCH DETECTED ---") logging.warning(f"Checkpoint '{os.path.basename(latest_checkpoint_path)}' was saved with different critical parameters:") for key, loaded_val, current_val in mismatched_keys: logging.warning(f" - {key}: Checkpoint='{loaded_val}', Current='{current_val}'") # Decide whether to proceed: raise error, warn, or try anyway # For now, just warn strongly. Loading might fail or lead to issues. logging.warning("Proceeding with loading, but results may be unexpected or errors may occur.") else: logging.warning("Checkpoint does not contain configuration info. Cannot check compatibility.") # --- End Config Check --- try: # Strict=False can sometimes help load partially, but hides potential issues model.load_state_dict(checkpoint['model'], strict=True) except RuntimeError as e: logging.error(f"Failed to load model state_dict: {e}") logging.error("This often happens due to architecture mismatch (check CONFIG) or corrupted checkpoint.") logging.error("Starting training from scratch.") return 0 # Cannot resume if model loading fails try: optimizer.load_state_dict(checkpoint['optimizer']) except ValueError as e: logging.warning(f"Could not load optimizer state_dict: {e}. Optimizer state will be reset.") # Reinitialize optimizer if state doesn't match? Or just proceed with current state. # Resetting optimizer state is safer if parameters changed. optimizer.state = defaultdict(dict) # Reset state logging.warning("Optimizer state reset.") except Exception as e: logging.error(f"Unexpected error loading optimizer state: {e}. Starting training from scratch.") return 0 start_step = checkpoint.get('step', 0) # Ensure step is non-negative, resume from next step start_step = max(0, start_step) + 1 if isinstance(start_step, int) else 0 logging.info(f"Checkpoint loaded successfully. Resuming from optimizer step {start_step}.") # Move optimizer state tensors to the correct device for state in optimizer.state.values(): for k, v in state.items(): if isinstance(v, torch.Tensor): try: state[k] = v.to(map_location) except Exception as e: logging.error(f"Failed to move optimizer tensor '{k}' to device '{map_location}': {e}") return start_step except FileNotFoundError: logging.info(f"No checkpoint directory '{self.checkpoint_dir}' or files found. Starting training from scratch.") return 0 except Exception as e: logging.error(f"Error loading checkpoint from '{self.checkpoint_dir}': {e}. Starting training from scratch.") # Clean up potentially partially loaded model/optimizer? # Re-initializing might be safer depending on where the error occurred. # For simplicity, we just return 0 here. return 0 # --- Training Function --- def train(): logging.info("Starting HROM training process on combined datasets (daily_dialog, empathetic_dialogues, blended_skill_talk, AlekseyKorshuk/persona-chat)...") # Corrected log message logging.info(f"Configuration: {CONFIG}") # --- Tokenizer Setup --- tokenizer_trainer = TokenizerTrainer() tokenizer_path = tokenizer_trainer.tokenizer_path if not os.path.exists(tokenizer_path): logging.info(f"Combined tokenizer '{CONFIG['tokenizer_name']}' not found. Training tokenizer...") try: # Pass trust_remote_code=True to load_dataset calls inside tokenizer training tokenizer_trainer.train(CONFIG["datasets"]) except Exception as e: logging.error(f"Failed during tokenizer training: {e}", exc_info=True) return # Cannot proceed without a tokenizer else: logging.info(f"Loading existing combined tokenizer from {tokenizer_path}") # Load the tokenizer instance *once* here for shared use try: tokenizer = tokenizer_trainer.get_tokenizer() # Update CONFIG with actual token IDs (useful for downstream) CONFIG['pad_token_id'] = tokenizer.token_to_id("") CONFIG['bos_token_id'] = tokenizer.token_to_id("") CONFIG['eos_token_id'] = tokenizer.token_to_id("") logging.info(f"Loaded tokenizer. Vocab size: {tokenizer.get_vocab_size()}. Special IDs: PAD={CONFIG['pad_token_id']}, BOS={CONFIG['bos_token_id']}, EOS={CONFIG['eos_token_id']}") except (FileNotFoundError, ValueError) as e: logging.error(f"Failed to load tokenizer: {e}. Cannot continue.") return # --- Model Initialization --- logging.info("Initializing HROM model...") # Ensure vocab_size in config matches tokenizer if CONFIG['vocab_size'] != tokenizer.get_vocab_size(): logging.warning(f"Config vocab_size ({CONFIG['vocab_size']}) differs from tokenizer vocab size ({tokenizer.get_vocab_size()}). Using tokenizer's size.") CONFIG['vocab_size'] = tokenizer.get_vocab_size() model = HROM() # --- Calculate and Log Model Parameters --- total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) logging.info(f"Model initialized. Total parameters: {total_params:,}") logging.info(f"Trainable parameters: {trainable_params:,}") logging.info(f"Parameters (Millions): Total={total_params/1e6:.2f}M, Trainable={trainable_params/1e6:.2f}M") # --- Dataset and DataLoader --- logging.info("Setting up combined dataset and dataloader...") try: logging.info("Pre-loading/caching datasets...") for ds_name in CONFIG["datasets"]: logging.info(f"Checking cache for '{ds_name}'...") try: # Load just the first example to trigger download/cache check _ = load_dataset(ds_name, split="train[:1]", download_mode="reuse_cache_if_exists", trust_remote_code=True) # Add trust_remote_code except Exception as e: # Log error but try to continue, main dataset loading will handle final error logging.error(f"Could not pre-check dataset '{ds_name}': {e}") logging.info("Dataset download/cache check presumed complete.") # Pass the already loaded tokenizer instance dataset = CombinedChatDataset(tokenizer) # Check if dataset is empty after processing if len(dataset) == 0: logging.error("Dataset is empty after processing all sources. Cannot train.") return dataloader = DataLoader( dataset, batch_size=CONFIG["batch_size"], collate_fn=CombinedChatDataset.collate_fn, # Use static method shuffle=True, # Adjust num_workers based on available cores, be conservative num_workers=min(4, os.cpu_count() // 2 if (os.cpu_count() and os.cpu_count() > 1) else 1), pin_memory=torch.cuda.is_available(), prefetch_factor=2 if torch.cuda.is_available() and os.cpu_count() and os.cpu_count() > 1 else None, drop_last=False # Keep last batch even if smaller ) except Exception as e: logging.error(f"Failed to initialize dataset/dataloader: {e}", exc_info=True) return # --- Trainer, Checkpoint, Safety --- logging.info("Initializing Trainer, Checkpoint Manager, and Safety Manager...") # Pass the loaded tokenizer instance trainer_obj = HROMTrainer(model, tokenizer) checkpoint_manager = CheckpointManager() # Uses CONFIG["checkpoint_dir"] safety = SafetyManager(model, tokenizer) # Pass the loaded tokenizer instance # --- Load Checkpoint --- start_optimizer_step = checkpoint_manager.load_latest(model, trainer_obj.optimizer) # Ensure model is on correct device after loading model.to(trainer_obj.device) # --- Training Loop --- logging.info(f"Starting training from optimizer step {start_optimizer_step}") optimizer_step = start_optimizer_step total_loss_accum = 0.0 # Calculate starting batch step based on loaded optimizer step and grad accum batch_step = optimizer_step * CONFIG["grad_accum_steps"] epochs_completed = batch_step // len(dataloader) if len(dataloader) > 0 else 0 start_epoch = epochs_completed # Start from the epoch corresponding to the loaded step # Estimate total steps (can be useful for LR scheduling if implementing decay) try: if len(dataloader) == 0: raise ValueError("DataLoader has zero length. Cannot estimate total steps.") total_optimizer_steps = (len(dataloader) * CONFIG["num_epochs"]) // CONFIG["grad_accum_steps"] logging.info(f"Estimated dataset size: {len(dataset)}") logging.info(f"Estimated batches per epoch: {len(dataloader)}") logging.info(f"Gradient Accumulation Steps: {CONFIG['grad_accum_steps']}") logging.info(f"Effective Batch Size: {CONFIG['batch_size'] * CONFIG['grad_accum_steps']}") logging.info(f"Target Epochs: {CONFIG['num_epochs']}") logging.info(f"Estimated total optimizer steps for {CONFIG['num_epochs']} epochs: {total_optimizer_steps}") except Exception as e: logging.warning(f"Could not accurately estimate dataloader length or total steps: {e}") total_optimizer_steps = -1 # Indicate unknown total steps model.train() # Ensure model is in training mode for epoch in range(start_epoch, CONFIG["num_epochs"]): logging.info(f"--- Starting Epoch {epoch+1}/{CONFIG['num_epochs']} ---") epoch_loss = 0.0 num_batches_in_epoch = 0 # Use enumerate starting from 1 for batch count if preferred for i, batch in enumerate(dataloader): # Check if batch is valid (collate_fn might return None) if batch is None: logging.warning(f"Skipping empty batch at step {i} in epoch {epoch+1}") continue # Forward and backward pass (scaled loss) loss = trainer_obj.train_step(batch) if loss is None or torch.isnan(torch.tensor(loss)) or torch.isinf(torch.tensor(loss)): logging.error(f"NaN, Inf, or None loss detected: {loss}. Epoch {epoch+1}, Batch {i}, Opt Step {optimizer_step}. Stopping.") # Try saving a 'nan_inf' checkpoint before exiting checkpoint_manager.save(model, trainer_obj.optimizer, f"{optimizer_step}_error") return total_loss_accum += loss epoch_loss += loss num_batches_in_epoch += 1 batch_step += 1 # Increment global batch counter (tracks batches processed) # Gradient Accumulation Check & Optimizer Step # Check if it's time to perform an optimizer step if batch_step % CONFIG["grad_accum_steps"] == 0: current_lr = trainer_obj.clip_and_step(optimizer_step) # Pass current opt step for LR schedule # Calculate average loss over accumulation steps for logging avg_loss = total_loss_accum / CONFIG["grad_accum_steps"] total_loss_accum = 0.0 # Reset loss accumulator # Logging if optimizer_step % CONFIG["debug_interval"] == 0: logging.info(f"Epoch {epoch+1} | Opt Step {optimizer_step} | Batch Step {batch_step} | Avg Loss: {avg_loss:.4f} | LR: {current_lr:.2e}") # Trigger debug generation less frequently or based on condition if optimizer_step % (CONFIG["debug_interval"] * 5) == 0: # e.g., every 5 debug intervals safety.debug_generation(" Hi there! How are you doing today?") # Use a generic debug prompt # Checkpointing if optimizer_step > 0 and optimizer_step % CONFIG["checkpoint_interval"] == 0: logging.info(f"Checkpoint interval reached at optimizer step {optimizer_step}.") checkpoint_manager.save(model, trainer_obj.optimizer, optimizer_step) # Optional: Run a generation check after saving checkpoint safety.debug_generation(" Hi! How are you?") optimizer_step += 1 # Increment optimizer step count *after* performing the step # --- End of Epoch --- avg_epoch_loss = epoch_loss / num_batches_in_epoch if num_batches_in_epoch > 0 else 0 logging.info(f"--- Finished Epoch {epoch+1}/{CONFIG['num_epochs']} | Average Epoch Loss: {avg_epoch_loss:.4f} ---") # Save checkpoint at the end of each epoch checkpoint_manager.save(model, trainer_obj.optimizer, f"epoch{epoch+1}_step{optimizer_step}") # Optionally run debug generation at end of epoch safety.debug_generation(" Hi! Whats up?") logging.info(f"Training finished after {CONFIG['num_epochs']} target epochs.") # Final save logging.info("Saving final model state...") checkpoint_manager.save(model, trainer_obj.optimizer, f"final_step{optimizer_step}") if __name__ == "__main__": # Ensures imports happen after setting the env var if script is run directly train()