train.py

# from peft import PeftModel, PeftConfig, get_peft_model
from datasets import load_dataset
from torch.utils.data import DataLoader
from peft import PeftModel, PeftConfig, get_peft_model
from utils.util import flatten_dict,shift_logits
from utils.data import get_bs17k_dataloader,get_llada_bs17k_dataloader,get_dataloader_by_config
from utils.model import get_model,get_llada,get_model_by_config
from utils.loss import compute_loss,compute_llada_loss,compute_normal_loss,compute_loss_by_config
from utils.generation import sample_tokens
# import dataloader

import os
import torch
import argparse
import torch.distributed as dist
from omegaconf import OmegaConf
from tqdm import tqdm
from accelerate import Accelerator
from accelerate.utils import ProjectConfiguration
import ipdb
os.environ['TOKENIZERS_PARALLELISM'] = 'false'

if os.environ.get("DEBUGPY", "0") == "1":
    import debugpy
    debugpy.listen(("0.0.0.0", 5678))
    print(f"Waiting for VS Code debugger attach (PID {os.getpid()})...", flush=True)
    debugpy.wait_for_client()


def get_accelerator(config, global_config):
    # Select experiment path based on config
    if hasattr(global_config, 'paths') and hasattr(global_config.paths, 'experiment'):
        root_path = global_config.paths.experiment
    else:
        root_path = config.root if hasattr(config, 'root') else '/tmp/experiment'
    
    output_dir = os.path.join(root_path, config.exp_name, config.output_dir)
    os.makedirs(output_dir, exist_ok=True)
    logging_dir = os.path.join(output_dir, config.logging_dir)
    project_config = ProjectConfiguration(project_dir=config.output_dir, logging_dir=logging_dir)
    accelerator = Accelerator(
        log_with=None if config.report_to == 'no' else config.report_to,
        mixed_precision=config.mixed_precision,
        project_config=project_config,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
    )

    return accelerator, output_dir

def main(args):
    config = OmegaConf.load(args.config)
    accelerator, output_dir = get_accelerator(config.train, config)
    # Use unified model and data loading functions
    denoiser, tokenizer = get_model_by_config(config)
    dataloader = get_dataloader_by_config(tokenizer, config.data, config)
    
    if config.train.decoder_resume_path is not None:
        ckpt = torch.load(config.train.decoder_resume_path, map_location='cpu', weights_only=True)
        if config.train.skipped_keys:
            ckpt = {k: v for k, v in ckpt.items() if k not in config.train.skipped_keys}
        m, u = denoiser.load_state_dict(ckpt, strict=False)
        if accelerator.is_main_process:
            print(f'model ckpt loaded from {config.train.decoder_resume_path}')

        # ckpt = torch.load(config.train.head_resume_path, map_location='cpu', weights_only=True)
        # if config.train.skipped_keys:
        #     ckpt = {k: v for k, v in ckpt.items() if k not in config.train.skipped_keys}
        # m, u = denoiser.lm_head.load_state_dict(ckpt, strict=False)
        # if accelerator.is_main_process:
        #     print(f'model ckpt loaded from {config.train.head_resume_path}')

    global_step = config.train.global_step if config.train.global_step is not None else 0
    params_to_learn = list(param for param in denoiser.parameters() if param.requires_grad)
    optimizer = torch.optim.AdamW(
        params_to_learn,
        lr           = config.train.lr,
        betas        = (0.9, 0.95),
        weight_decay = 5e-2,
        eps          = 1e-8,
    )
    
    denoiser, dataloader, optimizer = accelerator.prepare(denoiser, dataloader, optimizer)

    config.device_count = accelerator.num_processes
    if accelerator.is_main_process:
        accelerator.init_trackers(config.train.wandb_proj, config=flatten_dict(config))

    training_done = False
    epoch = 0
    progress_bar = tqdm(
        total   = config.train.num_iters,
        initial = global_step,
        desc    = 'Steps',
        disable = not accelerator.is_local_main_process,
    )

    if accelerator.is_main_process:
        print(f'Learnable parameters: {sum(p.numel() for p in params_to_learn if p.requires_grad) / 1e9} B')

    while not training_done:
        if accelerator.is_main_process:
            print(f'Epoch: {epoch}')
        for batch in dataloader:
            with accelerator.accumulate([denoiser]):
                denoiser.train()
                ipdb.set_trace()
                input_ids = batch['data']
                # print("input_ids",input_ids.dtype)
                question_length = batch['question_length']
                
                # Use unified loss function selection
                losses = compute_loss_by_config(
                    input_ids,
                    denoiser,
                    question_length,
                    block_size    = config.train.block_size,
                    mask_id       = config.denoiser.encoder.mask_id,
                    enable_shift  = config.train.enable_shift,
                    share_steps   = config.train.share_steps,
                    self_align    = config.train.self_align,
                    feature_align = config.train.feature_align,
                    self_step     = config.train.self_step,
                    eos_id        = tokenizer.eos_token_id,
                    config        = config
                )
                
                if config.train.share_steps > 1:
                    loss_tgt = losses['loss']
                    # loss_1 = losses['loss_1']
                    # loss_2 = losses['loss_2']
                else:
                    raise NotImplementedError
                torch.cuda.empty_cache()
                accelerator.backward(loss_tgt)
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(params_to_learn, 1.0)

                optimizer.step()
                optimizer.zero_grad()

            if accelerator.sync_gradients:
                global_step += 1
                progress_bar.update(1)
                logs = dict()
                loss_tgt = accelerator.gather(loss_tgt.detach()).mean().item()
                logs['loss'] = loss_tgt
                # if config.train.share_steps > 1:
                #     loss_1 = accelerator.gather(loss_1.detach()).mean().item()
                #     loss_2 = accelerator.gather(loss_2.detach()).mean().item()
                    # logs['loss_1'] = loss_1
                    # logs['loss_2'] = loss_2
                    
                accelerator.log(logs, step=global_step)
                progress_bar.set_postfix(**logs)

            if global_step > 0 and global_step % config.train.eval_every == 0 and accelerator.is_main_process:
                denoiser.eval();
                question = 'Henry made two stops during his 60-mile bike trip. He first stopped after 20 miles. His second stop was 15 miles before the end of the trip. How many miles did he travel between his first and second stops?'
                # prompt = tokenizer(question)['input_ids']
                # prompt = torch.tensor(prompt).to(accelerator.device).unsqueeze(0)
                messages = [
                    {"role": "user", "content": question}
                ]
                prompt = tokenizer.apply_chat_template(
                    messages, return_tensors="pt", return_dict=True, add_generation_prompt=True
                ).input_ids
                prompt = prompt.to(accelerator.device)

                mask_id = 151666
                gen_len = 512 - prompt.shape[1]
                temperature = 0.2
                top_p = 0.95

                x_t = torch.cat([prompt, torch.tensor([[mask_id]*gen_len]).to(accelerator.device)], dim=1)
                with torch.inference_mode():
                    for i in range(gen_len):
                        mask_index = (x_t == mask_id)
                        if i % 2 == 0:
                            z_t = denoiser.module.encoder(x_t, output_hidden_states=True).hidden_states[-1]
                            hidden_state = denoiser.module.decoder(x_t, z_t)
                            logits = denoiser.module.encoder.lm_head(hidden_state)
                        else:
                            hidden_state = denoiser.module.decoder(x_t, z_t)
                            logits = denoiser.module.lm_head(hidden_state)

                        if config.train.enable_shift:
                            logits = shift_logits(logits)

                        mask_logits = logits[mask_index]
                        confidence, x0 = sample_tokens(mask_logits, temperature, top_p=top_p, top_k=None, neg_entropy=True)

                        number_transfer_tokens = 1
                        _, transfer_index = torch.topk(confidence, number_transfer_tokens)
                        x0_ = torch.zeros_like(x0, device=accelerator.device, dtype=torch.long) + mask_id
                        x0_[transfer_index] = x0[transfer_index].clone()
                        x_t[mask_index] = x0_

                answer = tokenizer.batch_decode(x_t[:, prompt.shape[1]:], skip_special_tokens=True)[0]
                print(answer)

            accelerator.wait_for_everyone()

            if global_step > 0 and global_step % config.train.save_every == 0 and accelerator.is_main_process:
                denoiser.eval()
                decoder_state_dict = accelerator.unwrap_model(denoiser).save_pretrained(os.path.join(output_dir, f"Decoder-{config.train.exp_name}-{global_step // 1000}k"))
                # lmhead_state_dict = accelerator.unwrap_model(denoiser).lm_head.state_dict()
                # torch.save(lmhead_state_dict, os.path.join(output_dir, f"LMhead-{config.train.exp_name}-{global_step // 1000}k"))
            accelerator.wait_for_everyone()
            if global_step >= config.train.num_iters:
                training_done = True
                break
        epoch += 1
    accelerator.end_training()
    if dist.is_initialized():
        dist.destroy_process_group()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='config/dream.yaml')
    args = parser.parse_args()
    main(args)