train_hope.py

from __future__ import annotations

from collections import deque
import hydra
import json
import math
import random
import time
import torch
import torch.nn as nn
import multiprocessing as mp
from omegaconf import DictConfig
from queue import Empty
from transformers import AutoTokenizer, 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.remote_data import RemoteDataClient
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 (
    _check_state_recursive,
    _log_nan_debug,
    _check_tensor,
    setup_nan_debug,
    setup_perf_logger,
    setup_clipping_debug,
)
import gc
from datasets import load_dataset

MAX_TOKENIZER_CHUNK = 131072


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:
        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 tokenize_texts_to_sample_items(
    tokenizer,
    texts: list[str],
    *,
    max_tokenizer_chunk: int = MAX_TOKENIZER_CHUNK,
) -> list[dict[str, list[int]]]:
    if max_tokenizer_chunk <= 0:
        raise ValueError(f"max_tokenizer_chunk must be positive, got {max_tokenizer_chunk}")
    if not texts:
        return []

    tokenized_batch = tokenizer(
        texts,
        add_special_tokens=False,
        truncation=False,
        return_attention_mask=False,
    )

    sample_items: list[dict[str, list[int]]] = []
    chunked_count = 0
    for text, input_ids in zip(texts, tokenized_batch["input_ids"], strict=True):
        if len(input_ids) <= max_tokenizer_chunk:
            sample_items.append({"input_ids": input_ids, "raw_text": text})
            continue

        chunked_count += 1
        for start_idx in range(0, len(input_ids), max_tokenizer_chunk):
            chunk_ids = input_ids[start_idx:start_idx + max_tokenizer_chunk]
            if chunk_ids:
                sample_items.append({"input_ids": chunk_ids, "raw_text": text})

    if chunked_count > 0:
        print(
            "[DataGen] Chunked long texts "
            f"texts={chunked_count} max_tokenizer_chunk={max_tokenizer_chunk} "
            f"output_samples={len(sample_items)}"
        )
    return sample_items


def data_generator(
    dataset_name: str,
    data_split: str,
    data_server_ip: str | None,
    data_server_port: int | None,
    data_server_route_prefix: str,
    data_server_timeout: float,
    data_server_seed: int,
    data_server_shuffle_buffer_size: int,
    producer_batch_size: int,
    batch_seq_len: int,
    pad_token_id: int,
    step_size: int,
    y_context_window: int,
    y_encoder_model_name: str,
    y_encoder_batch_size: int,
    device: str,
):
    """Generator that streams, tokenizes, materializes samples and yields one data item at a time."""
    print("[DataGen] Initializing encoders and dataset...")
    y_encoder_model = AutoModel.from_pretrained(y_encoder_model_name, trust_remote_code=True)
    y_encoder_model.to(device).eval()
    embedding_layer = get_embedding_layer(y_encoder_model)
    tokenizer = AutoTokenizer.from_pretrained(y_encoder_model_name, use_fast=True)

    stream_batch_manager = StreamingBatchManager(
        batch_seq_len=batch_seq_len,
        pad_token_id=pad_token_id,
    )

    use_remote_data_server = bool(data_server_ip and data_server_port)
    train_stream = None
    remote_client = None
    try:
        if use_remote_data_server:
            remote_client = RemoteDataClient(
                host=data_server_ip,
                port=int(data_server_port),
                route_prefix=data_server_route_prefix,
                timeout=float(data_server_timeout),
                seed=int(data_server_seed),
                shuffle_buffer_size=int(data_server_shuffle_buffer_size),
            )
            health = remote_client.healthcheck()
            session = remote_client.open_session()
            print(
                "[DataGen] Connected to remote data server "
                f"host={data_server_ip} port={data_server_port} "
                f"records={session.get('total_records')} health={health.get('status')}"
            )
        else:
            dataset = load_dataset(dataset_name, split=data_split, streaming=True)
            train_stream = dataset.batch(producer_batch_size)

        print("[DataGen] Starting to generate data...")
        step = 0
        while True:
            if use_remote_data_server:
                assert remote_client is not None
                batch = remote_client.next_batch(producer_batch_size)
                texts = batch.get("texts", [])
                if not texts:
                    print("[DataGen] Remote data server exhausted")
                    break
            else:
                assert train_stream is not None
                batch = next(train_stream, None)
                if batch is None:
                    break
                texts = batch.get("text", [])

            sample_items = tokenize_texts_to_sample_items(
                tokenizer,
                texts,
            )
            if not sample_items:
                step += 1
                continue

            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,
                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"]
                raw_text_batch = group.get("raw_text_batch", [None] * len(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]):
                    yield {
                        "step": step,
                        "aligned_seq_len": aligned_seq_len,
                        "raw_input_ids": raw_input_ids_batch[row_idx],
                        "raw_text": raw_text_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],
                    }
                    produced_count += 1
                    if produced_count % 64 == 0:
                        print(
                            f"[DataGen] step={step} samples={produced_count}"
                        )

            del materialized_batch
            gc.collect()
            if str(device).startswith("cuda"):
                torch.cuda.empty_cache()
            step += 1

        print("[DataGen] Finished generating all data")
    finally:
        if remote_client is not None:
            remote_client.close()


def producer_worker(
    data_queue,
    control_queue,
    dataset_name: str,
    data_split: str,
    data_server_ip: str | None,
    data_server_port: int | None,
    data_server_route_prefix: str,
    data_server_timeout: float,
    data_server_seed: int,
    data_server_shuffle_buffer_size: int,
    producer_batch_size: int,
    train_batch_size: int,
    batch_seq_len: int,
    step_size: int,
    y_context_window: int,
    y_encoder_model_name: str,
    y_encoder_batch_size: int,
    producer_device: str,
) -> None:
    """Background producer that materializes samples and pushes them to a queue."""
    rng = random.Random(data_server_seed)

    def emit_ready_batches(pending_samples: list[dict], *, force: bool = False) -> None:
        while len(pending_samples) >= train_batch_size or (force and pending_samples):
            batch_len = train_batch_size if len(pending_samples) >= train_batch_size else len(pending_samples)
            batch_samples = pending_samples[:batch_len]
            del pending_samples[:batch_len]
            data_queue.put(batch_samples)

    def signal_done() -> None:
        try:
            data_queue.put(None, timeout=1)
        except Exception:
            pass

    remote_client = None
    try:
        print("[Producer] Initializing encoders and dataset...")
        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)
        tokenizer = AutoTokenizer.from_pretrained(y_encoder_model_name, use_fast=True)
        pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id

        stream_batch_manager = StreamingBatchManager(
            batch_seq_len=batch_seq_len,
            pad_token_id=pad_token_id,
        )

        use_remote_data_server = bool(data_server_ip and data_server_port)
        train_stream = None
        if use_remote_data_server:
            remote_client = RemoteDataClient(
                host=data_server_ip,
                port=int(data_server_port),
                route_prefix=data_server_route_prefix,
                timeout=float(data_server_timeout),
                seed=int(data_server_seed),
                shuffle_buffer_size=int(data_server_shuffle_buffer_size),
            )
            health = remote_client.healthcheck()
            session = remote_client.open_session()
            print(
                "[Producer] Connected to remote data server "
                f"host={data_server_ip} port={data_server_port} "
                f"records={session.get('total_records')} health={health.get('status')}"
            )
        else:
            dataset = load_dataset(dataset_name, split=data_split, streaming=True)
            train_stream = dataset.batch(producer_batch_size)

        print("[Producer] Starting to generate data...")
        step = 0
        pending_samples: list[dict] = []
        while True:
            try:
                msg = control_queue.get_nowait()
                if msg == "STOP":
                    print("[Producer] Received stop signal")
                    break
            except Empty:
                pass

            if use_remote_data_server:
                assert remote_client is not None
                batch = remote_client.next_batch(producer_batch_size)
                texts = batch.get("texts", [])
                if not texts:
                    print("[Producer] Remote data server exhausted")
                    break
            else:
                assert train_stream is not None
                batch = next(train_stream, None)
                if batch is None:
                    break
                texts = batch.get("text", [])

            sample_items = tokenize_texts_to_sample_items(
                tokenizer,
                texts,
            )
            if not sample_items:
                step += 1
                continue

            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"]
                raw_text_batch = group.get("raw_text_batch", [None] * len(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]):
                    try:
                        msg = control_queue.get_nowait()
                        if msg == "STOP":
                            print("[Producer] Received stop signal")
                            signal_done()
                            return
                    except Empty:
                        pass

                    pending_samples.append(
                        {
                            "producer_step": step,
                            "aligned_seq_len": aligned_seq_len,
                            "raw_input_ids": raw_input_ids_batch[row_idx],
                            "raw_text": raw_text_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],
                        }
                    )
                    produced_count += 1
                    if len(pending_samples) >= producer_batch_size:
                        rng.shuffle(pending_samples)
                        emit_ready_batches(pending_samples)
                    if produced_count % 64 == 0:
                        queue_size = "unknown"
                        if hasattr(data_queue, "qsize"):
                            try:
                                queue_size = data_queue.qsize()
                            except NotImplementedError:
                                pass
                        print(
                            f"[Producer] step={step} samples={produced_count} queue_size~{queue_size}"
                        )

            del materialized_batch
            gc.collect()
            if str(producer_device).startswith("cuda"):
                torch.cuda.empty_cache()
            step += 1

        if pending_samples:
            rng.shuffle(pending_samples)
            emit_ready_batches(pending_samples, force=True)
        signal_done()
        print("[Producer] Finished generating all data")
    except Exception as exc:
        print(f"[Producer] Error: {exc}")
        import traceback
        traceback.print_exc()
        signal_done()
    finally:
        if remote_client is not None:
            remote_client.close()


@hydra.main(config_path="configs", config_name="hope_tiny", version_base=None)
def main(cfg: DictConfig) -> None:
    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)

    model.is_training = True
    step_size = model.min_step_size
    y_context_window = cfg.train.y_context_window

    dataset_name = cfg.data.source
    print(f"Loading dataset: {dataset_name}")
    data_split = cfg.data.split
    data_server_ip = cfg.data.get("data_server_ip", None)
    data_server_port = cfg.data.get("data_server_port", None)
    data_server_route_prefix = cfg.data.get("data_server_route_prefix", "")
    data_server_timeout = cfg.data.get("data_server_timeout", 60.0) 
    data_server_seed = cfg.data.get("data_server_seed", random.randint(0, 1000000))
    data_server_shuffle_buffer_size = cfg.data.get("data_server_shuffle_buffer_size", 8192)

    y_encoder_model_name = cfg.train.y_encoder_model_name

    batch_size = cfg.train.get("batch_size", 1)
    y_encoder_batch_size = cfg.train.get("y_encoder_batch_size", 1)
    producer_batch_size = cfg.train.get("producer_batch_size", batch_size)
    producer_device = cfg.train.get("producer_device", device)

    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", cfg.train.get("steps", None))

    mp.set_start_method("spawn", force=True)
    queue_buffer_size = cfg.train.get("queue_buffer_size", batch_size * 4)
    local_buffer_size = cfg.train.get("local_buffer_size", batch_size * 2)
    local_buffer_threshold = cfg.train.get("local_buffer_threshold", batch_size)
    if queue_buffer_size <= 0:
        raise ValueError(f"queue_buffer_size must be positive, got {queue_buffer_size}")
    if local_buffer_size <= 0:
        raise ValueError(f"local_buffer_size must be positive, got {local_buffer_size}")
    if local_buffer_threshold <= 0:
        raise ValueError(f"local_buffer_threshold must be positive, got {local_buffer_threshold}")
    if local_buffer_threshold > local_buffer_size:
        raise ValueError(
            "local_buffer_threshold must be less than or equal to local_buffer_size, "
            f"got {local_buffer_threshold} and {local_buffer_size}"
        )

    data_queue = mp.Queue(maxsize=queue_buffer_size)
    control_queue = mp.Queue()
    producer_process = mp.Process(
        target=producer_worker,
        args=(
            data_queue,
            control_queue,
            dataset_name,
            data_split,
            data_server_ip,
            data_server_port,
            data_server_route_prefix,
            data_server_timeout,
            data_server_seed,
            data_server_shuffle_buffer_size,
            producer_batch_size,
            batch_size,
            batch_seq_len,
            step_size,
            y_context_window,
            y_encoder_model_name,
            y_encoder_batch_size,
            producer_device,
        ),
    )
    producer_process.start()
    print(
        "[Main] Producer process started "
        f"(queue_buffer_size={queue_buffer_size}, "
        f"local_buffer_size={local_buffer_size}, "
        f"local_buffer_threshold={local_buffer_threshold})"
    )

    outer_step = 0
    cached_outer_grads_sum: dict[str, torch.Tensor] = {}
    cached_outer_grads_count: dict[str, int] = {}
    cached_outer_block_update_steps: dict[str, int] = {}
    last_block_need_update_dict: dict[str, bool] = None
    active_samples: list[dict] = []
    local_buffer: deque[dict] = deque()
    producer_done = False

    def try_fill_local_buffer(*, block: bool, min_count: int = 0) -> None:
        nonlocal producer_done
        while len(local_buffer) < local_buffer_size and not producer_done:
            should_block = block and len(local_buffer) < min_count
            try:
                if should_block:
                    item = data_queue.get(timeout=data_server_timeout)
                else:
                    item = data_queue.get_nowait()
            except Empty:
                break

            if item is None:
                producer_done = True
                break
            if not isinstance(item, list):
                raise TypeError(f"Expected queued batch list, got {type(item).__name__}")
            local_buffer.extend(item)

    def refill_active_samples(*, block: bool, min_count: int = 1) -> None:
        if len(active_samples) >= min_count:
            return
        try_fill_local_buffer(block=block, min_count=min_count)
        while local_buffer and len(active_samples) < min_count:
            active_samples.append(local_buffer.popleft())

    try:
        refill_active_samples(block=True, min_count=batch_size)

        if not active_samples:
            print("[Main] No data available")
            return

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

        while total_steps is None or outer_step < total_steps:
            if not active_samples:
                refill_active_samples(block=True, min_count=1)
                if not active_samples:
                    print("[Main] No active samples available, waiting for data")
                    if producer_done:
                        print("[Main] Producer exhausted and no active samples remain")
                        break
                    continue

            if len(local_buffer) < local_buffer_threshold and not producer_done:
                try_fill_local_buffer(block=False)

            start_time = time.time()
            raw_input_ids_batch = [sample["raw_input_ids"] for sample in active_samples]
            raw_text_batch = [sample.get("raw_text") 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, last_block_need_update_dict = 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,
                last_block_need_update_dict=last_block_need_update_dict,
            )

            _has_nan_loss = any(
                any(math.isnan(v) or math.isinf(v) for v in row)
                for row in outer_seq_loss
            )
            if _has_nan_loss:
                _log_nan_debug(
                    f"outer_step={outer_step} *** outer_seq_loss contains NaN/Inf! "
                    f"losses={outer_seq_loss}"
                )
                _state_bad = _check_state_recursive(f"outer_step={outer_step} stream_state", stream_state)
                if _state_bad:
                    _log_nan_debug(f"outer_step={outer_step} *** {_state_bad} state tensors have NaN/Inf")
                _log_nan_debug(f"outer_step={outer_step} x stats: min={x.min().item():.4g} max={x.max().item():.4g} has_nan={torch.isnan(x).any().item()}")
                _log_nan_debug(f"outer_step={outer_step} y stats: min={y.min().item():.4g} max={y.max().item():.4g} has_nan={torch.isnan(y).any().item()}")

                for _pname, _pval in model.pp_local_trainable_named_parameters().items():
                    _check_tensor(f"outer_step={outer_step} model_param[{_pname}]", _pval)

            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)} "
                f"window={producer_buffer_size} local_buffer={len(local_buffer)} "
                f"elapsed={elapsed_time:.3f}s losses={outer_batch_loss}"
            )

            payload = []
            for seq_loss, input_ids, raw_text in zip(outer_seq_loss_trimmed, raw_input_ids_batch, raw_text_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}; {json.dumps(raw_text, ensure_ascii=False)}"
                )
            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"][:, batch_seq_len:]
                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:
                try_fill_local_buffer(block=True, min_count=len(finished_indices))

                replacement_samples = []
                replacement_indices = []
                for idx in finished_indices:
                    if not local_buffer:
                        break
                    replacement_samples.append(local_buffer.popleft())
                    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)
                    unfilled_set = set(unfilled_indices)
                    keep_indices = [i for i in range(current_batch_size) if i not in unfilled_set]
                    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)} (data exhausted)")

            if not active_samples and producer_done:
                print("[Main] No active samples remain after replacement")
                del x, y, seq_mask
                gc.collect()
                if str(device).startswith("cuda"):
                    torch.cuda.empty_cache()
                break

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

            outer_step += 1
    finally:
        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()