train_hope_debug.py

from __future__ import annotations

import hydra
import time
import torch
import torch.nn as nn
from omegaconf import DictConfig
from transformers import AutoModel
from nested_learning_pytorch.training import build_model_from_cfg, unwrap_config
from nested_learning_pytorch.manager import StreamingBatchManager, materialize_samples
from nested_learning_pytorch.utils import count_params, maybe_save_checkpoint, keep_state_rows
from nested_learning_pytorch.logging import init_logger
from nested_learning_pytorch.model import setup_clipping_debug, setup_nan_debug, setup_perf_logger
import gc
import multiprocessing as mp
from queue import Empty


def get_embedding_layer(model: nn.Module) -> nn.Module:
    if hasattr(model, 'embed_tokens'):
        embedding_layer = model.embed_tokens
    elif hasattr(model, 'transformers_model') and hasattr(model.transformers_model, 'embed_tokens'):
        embedding_layer = model.transformers_model.embed_tokens
    elif hasattr(model, 'model') and hasattr(model.model, 'embed_tokens'):
        embedding_layer = model.model.embed_tokens
    elif hasattr(model, 'transformer') and hasattr(model.transformer, 'wte'):
        embedding_layer = model.transformer.wte
    else:
        # Fallback: try to find embedding layer
        print("Model structure:")
        for name, module in model.named_modules():
            if 'embed' in name.lower():
                print(f"  {name}: {type(module)}")
        raise ValueError("Could not find embedding layer in model")
    return embedding_layer

def producer_worker(data_queue, control_queue, train_stream_gen_func, train_stream_args,
                    batch_seq_len, pad_token_id, step_size,
                    y_context_window, y_encoder_model_name, y_encoder_batch_size,
                    producer_device):
    """Producer process: batch materialize samples then push one sample per queue item."""
    try:
        print("[Producer] Initializing y encoder and embedding layer...")
        y_encoder_model = AutoModel.from_pretrained(y_encoder_model_name, trust_remote_code=True)
        y_encoder_model.to(producer_device).eval()
        embedding_layer = get_embedding_layer(y_encoder_model)
        stream_batch_manager = StreamingBatchManager(
            batch_seq_len=batch_seq_len,
            pad_token_id=pad_token_id,
        )

        train_stream = train_stream_gen_func(*train_stream_args)

        print("[Producer] Starting to produce data...")
        step = 0
        for batch in train_stream:
            for sample_idx in range(len(batch["input_ids"])):
                try:
                    msg = control_queue.get_nowait()
                    if msg == "STOP":
                        print("[Producer] Received stop signal")
                        return
                except Empty:
                    pass

            sample_items = [{"input_ids": ids} for ids in batch["input_ids"]]
            materialized_batch = materialize_samples(
                sample_items,
                stream_batch_manager,
                embedding_layer,
                step_size,
                y_context_window,
                pad_token_id,
                y_encoder_model,
                y_encoder_batch_size,
                producer_device,
            )

            produced_count = 0
            for group in materialized_batch:
                aligned_seq_len = group["aligned_seq_len"]
                raw_input_ids_batch = group["raw_input_ids_batch"]
                x_padded = group["x_padded"]
                y_padded = group["y_padded"]
                seq_mask_padded_batch = group["seq_mask_padded_batch"]

                for row_idx in range(x_padded.shape[0]):
                    data_item = {
                        "producer_step": step,
                        "aligned_seq_len": aligned_seq_len,
                        "raw_input_ids": raw_input_ids_batch[row_idx],
                        "x_padded": x_padded[row_idx:row_idx + 1],
                        "y_padded": y_padded[row_idx:row_idx + 1],
                        "seq_mask_padded_batch": seq_mask_padded_batch[row_idx:row_idx + 1],
                    }
                    data_queue.put(data_item)
                    produced_count += 1
                    if produced_count % 64 == 0:
                        print(f"[Producer] Produced step={step} samples={produced_count}, queue size: ~{data_queue.qsize()}")
            del materialized_batch
            gc.collect()
            if str(producer_device).startswith('cuda'):
                torch.cuda.empty_cache()
            step += 1

        data_queue.put(None)
        print("[Producer] Finished producing all data")
        
    except Exception as e:
        print(f"[Producer] Error: {e}")
        import traceback
        traceback.print_exc()
        data_queue.put(None)


def create_train_stream(log_file_path, batch_size):
    """Generator that reads train_hope.log and yields batches with input_ids"""
    import ast
    with open(log_file_path, 'r') as f:
        batch_input_ids = []
        for line in f:
            # Parse line: "loss; [input_ids]"
            if '; ' in line:
                _, _, input_ids_str = line.strip().split('; ', 2)
                # Convert string representation of list to actual list
                input_ids = ast.literal_eval(input_ids_str)
                batch_input_ids.append(input_ids)
                
                # Yield batch when we have enough samples
                if len(batch_input_ids) >= batch_size:
                    yield {"input_ids": batch_input_ids}
                    batch_input_ids = []
        
        # Yield remaining samples if any
        if batch_input_ids:
            yield {"input_ids": batch_input_ids}


@hydra.main(config_path="configs", config_name="hope_tiny", version_base=None)
def main(cfg: DictConfig) -> None:
    # build model
    cfg = unwrap_config(cfg)
    device = cfg.train.device
    
    print(f"Building model from config: {cfg.model}")
    model = build_model_from_cfg(cfg.model, cfg.train)
    # maybe_save_checkpoint(cfg, model, step=-1, total_steps=100)
    model.load_checkpoint(cfg.train.get("load_checkpoint_path", None))
    # count_params(model)
    
    logger = init_logger(getattr(cfg, "logging", None), cfg)

    _log_cfg = getattr(cfg, "logging", None)
    _nan_debug = bool(_log_cfg.get("nan_debug", True)) if _log_cfg is not None else True
    _perf_debug = bool(_log_cfg.get("perf_debug", False)) if _log_cfg is not None else False
    _clipping_debug = bool(_log_cfg.get("clipping_debug", False)) if _log_cfg is not None else False
    setup_nan_debug(enabled=_nan_debug)
    setup_perf_logger(enabled=_perf_debug)
    setup_clipping_debug(enabled=_clipping_debug)
    if _perf_debug:
        print("[Main] Performance debug logging enabled → train_hope_debug.log")

    model.is_training = True
    step_size = model.min_step_size
    y_context_window = cfg.train.y_context_window
    outer_step = model.step
    
    # load dataset
    dataset_name = cfg.data.source
    print(f"Loading dataset: {dataset_name}")
    
    y_encoder_model_name = cfg.train.y_encoder_model_name
    
    # batch dataset using default method
    batch_size = cfg.train.get("batch_size", 1)
    y_encoder_batch_size = cfg.train.get("y_encoder_batch_size", 1)
    
    # Get pad token id from tokenizer (use eos_token as pad if pad_token not set)
    pad_token_id = 151643
    
    print("Training...")
    batch_seq_len = cfg.train.batch_seq_len
    if batch_seq_len % step_size != 0:
        raise ValueError(f"batch_seq_len ({batch_seq_len}) must be divisible by step_size ({step_size})")
    producer_buffer_size = cfg.train.get("producer_buffer_size", batch_seq_len)
    if producer_buffer_size <= 0:
        raise ValueError(f"producer_buffer_size must be positive, got {producer_buffer_size}")
    per_outer_optimize_steps = cfg.train.get("per_outer_optimize_steps", 1)
    if per_outer_optimize_steps <= 0:
        raise ValueError(f"per_outer_optimize_steps must be positive, got {per_outer_optimize_steps}")
    per_outer_calculation_steps = cfg.train.get("per_outer_calculation_steps", 1)
    if per_outer_calculation_steps <= 0:
        raise ValueError(
            f"per_outer_calculation_steps must be positive, got {per_outer_calculation_steps}"
        )
    if per_outer_optimize_steps % per_outer_calculation_steps != 0:
        raise ValueError(
            "per_outer_optimize_steps must be divisible by per_outer_calculation_steps, "
            f"got {per_outer_optimize_steps} and {per_outer_calculation_steps}"
        )
    total_steps = cfg.train.get("total_steps", None)
    producer_batch_size = cfg.train.get("producer_batch_size", batch_size)
    producer_device = cfg.train.get("producer_device", device)
    
    # Set multiprocessing start method
    mp.set_start_method('spawn', force=True)
    
    # Create data queue and control queue
    buffer_size = cfg.train.get("buffer_size", 100)  # Queue buffer size
    data_queue = mp.Queue(maxsize=buffer_size)
    control_queue = mp.Queue()
    
    # Start producer process
    producer_process = mp.Process(
        target=producer_worker,
        args=(
            data_queue, control_queue,
            create_train_stream, ('train_hope.log', producer_batch_size),
            batch_seq_len, pad_token_id, step_size,
            y_context_window, y_encoder_model_name, y_encoder_batch_size,
            producer_device,
        )
    )
    producer_process.start()
    print("[Main] Producer process started")
    
    try:
        producer_done = False
        cached_outer_grads_sum: dict[str, torch.Tensor] = {}
        cached_outer_grads_count: dict[str, int] = {}
        cached_outer_block_update_steps: dict[str, int] = {}
        active_samples: list[dict] = []

        while len(active_samples) < batch_size:
            data_item = data_queue.get()
            if data_item is None:
                producer_done = True
                break
            active_samples.append(data_item)

        if not active_samples:
            print("[Main] No data received from producer")
            return

        stream_state = model.init_state(device=device, batch_size=len(active_samples))
        y_steps_per_outer = batch_seq_len // step_size

        while active_samples and (total_steps is None or outer_step < total_steps):
            start_time = time.time()
            raw_input_ids_batch = [sample["raw_input_ids"] for sample in active_samples]
            x_batch_chunks = []
            y_batch_chunks = []
            mask_batch_chunks = []
            is_batch_finished = []

            for sample in active_samples:
                if sample["x_padded"].shape[1] < producer_buffer_size:
                    raise ValueError(
                        "Remaining sequence length is smaller than producer_buffer_size; "
                        "please ensure alignment config matches producer_buffer_size."
                    )
                x_batch_chunks.append(sample["x_padded"][:, :batch_seq_len])
                y_batch_chunks.append(sample["y_padded"][:, :batch_seq_len])
                mask_batch_chunks.append(sample["seq_mask_padded_batch"][:, :y_steps_per_outer])
                is_batch_finished.append(sample["x_padded"].shape[1] <= batch_seq_len)

            x = torch.cat(x_batch_chunks, dim=0).to(device)
            y = torch.cat(y_batch_chunks, dim=0).to(device)
            seq_mask = torch.cat(mask_batch_chunks, dim=0).to(device)
            batch_finished_mask = torch.tensor(is_batch_finished, dtype=torch.float32, device=device)

            outer_seq_loss, stream_state = model.forward_inner_loop(
                x=x,
                y=y,
                seq_mask=seq_mask,
                batch_finished_mask=batch_finished_mask,
                state=stream_state,
                per_outer_optimize_steps=per_outer_optimize_steps,
                per_outer_calculation_steps=per_outer_calculation_steps,
                cached_outer_grads_sum=cached_outer_grads_sum,
                cached_outer_grads_count=cached_outer_grads_count,
                cached_outer_block_update_steps=cached_outer_block_update_steps,
            )
            maybe_save_checkpoint(cfg, model, step=outer_step, total_steps=total_steps)

            seq_mask_list = seq_mask.clone().cpu().tolist()
            outer_seq_loss_trimmed = [
                seq_loss[:int(sum(mask_row))]
                for seq_loss, mask_row in zip(outer_seq_loss, seq_mask_list, strict=True)
            ]
            outer_batch_loss = [round(sum(seq_loss) / len(seq_loss), 3) for seq_loss in outer_seq_loss_trimmed]

            elapsed_time = time.time() - start_time
            print(f"[Main] outer_step={outer_step} batch={len(active_samples)} window={producer_buffer_size} elapsed={elapsed_time:.3f}s losses={outer_batch_loss}")

            payload = []
            for seq_loss, input_ids in zip(outer_seq_loss_trimmed, raw_input_ids_batch, strict=True):
                payload.append(
                    f'{round(sum(seq_loss) / len(seq_loss), 3)}; '
                    f'{[round(loss, 3) for loss in seq_loss]}; {input_ids}'
                )
            logger.log(payload, step=outer_step)

            for sample in active_samples:
                sample["x_padded"] = sample["x_padded"][:, batch_seq_len:]
                sample["y_padded"] = sample["y_padded"][:, y_steps_per_outer:]
                sample["seq_mask_padded_batch"] = sample["seq_mask_padded_batch"][:, y_steps_per_outer:]

            finished_indices = [
                idx for idx, sample in enumerate(active_samples)
                if sample["x_padded"].shape[1] == 0
            ]
            if finished_indices:
                replacement_samples = []
                replacement_indices = []
                for idx in finished_indices:
                    if producer_done:
                        break
                    item = data_queue.get()
                    if item is None:
                        producer_done = True
                        break
                    replacement_samples.append(item)
                    replacement_indices.append(idx)

                for idx, replacement in zip(replacement_indices, replacement_samples, strict=True):
                    active_samples[idx] = replacement

                unfilled_indices = sorted(set(finished_indices) - set(replacement_indices), reverse=True)
                if unfilled_indices:
                    current_batch_size = len(active_samples)
                    keep_indices = [i for i in range(current_batch_size) if i not in set(unfilled_indices)]
                    active_samples = [active_samples[i] for i in keep_indices]
                    stream_state = keep_state_rows(stream_state, keep_indices, current_batch_size=current_batch_size)
                    print(f"[Main] Shrunk active batch to {len(active_samples)} (producer exhausted)")

            del x, y, seq_mask
            gc.collect()
            if device.startswith('cuda'):
                torch.cuda.empty_cache()

            outer_step += 1

    finally:
        # Cleanup producer process
        print("[Main] Cleaning up producer process...")
        control_queue.put("STOP")
        producer_process.join(timeout=5)
        if producer_process.is_alive():
            print("[Main] Force terminating producer process...")
            producer_process.terminate()
            producer_process.join()
        print("[Main] Producer process stopped")


if __name__ == "__main__":
    main()