language_graph_pretrain.py

from data import ( 
    MultiModalDataset, 
    collate_fn,
    convert_batch_to_unirtl_input_format,
    compute_class_counts
)
import torch, json
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader, DistributedSampler, random_split
from torch import nn
import torch_geometric
from transformers import AutoTokenizer, BertConfig, RobertaConfig, get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup
from models.bert_new import BertForPreTrainingWithMlmAndNspWithGraph
from models.roberta_new import RobertaForMaskedLMWithGraph
from models.unirtl_model import UniRTL, UniRTLConfig
from models.graph_aware_tokenizer.cdfg_emb import load_model
import os, argparse
import logging
from datetime import datetime
import argparse

def main():
    torch.multiprocessing.set_sharing_strategy('file_system')
    parser = argparse.ArgumentParser(description="UniRTL Training Script")
    parser.add_argument("--per_device_batch_size", type=int, default=4, help="Per GPU batch size")
    parser.add_argument("--bert_model_id", type=str, default="bert-base-uncased", help="Pretrained BERT model identifier")
    parser.add_argument("--model_type", type=str, default="bert", help="type of bert model [bert, roberta]")
    parser.add_argument("--cdfgnn_model_path", type=str, default="language_pretrain/interface/ckpts/32_class_best_model", help="Path to CdfgNN model")
    parser.add_argument("--max_text_len", type=int, default=512, help="Maximum text sequence length")
    parser.add_argument("--max_graph_len", type=int, default=512, help="Maximum graph node number")
    parser.add_argument("--text_mask_ratio", type=float, default=0.2, help="Mask how much of the description and code tokens")
    parser.add_argument("--num_node_classes", type=int, default=32, help="num of label classes for graph node type prediction")
    parser.add_argument("--graph_mask_prob", type=float, default=0.15, help="graph node mask probability")
    parser.add_argument("--epochs", type=int, default=1, help="Number of training epochs")
    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--save_path", type=str, default="first_dist_run", help="Path to save checkpoints and logs")
    parser.add_argument("--log_steps", type=int, default=100, help="Log every N steps")
    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping max norm")
    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay for optimizer")
    parser.add_argument("--val_split_ratio", type=float, default=0.1, help="Validation set ratio")
    parser.add_argument("--lr_scheduler_type", type=str, default="linear", choices=["linear", "cosine"], help="Learning rate scheduler type")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of steps for gradient accumulation")
    parser.add_argument("--graph_loss_type", type=str, choices=["ce", "cb_focal"], default="cb_focal", help="loss function to use: ce (CrossEntropy), cb_focal (ClassBalancedFocalLoss)")
    args = parser.parse_args()

    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(device)

    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    is_main_process = rank == 0
    if is_main_process:
        os.makedirs(args.save_path, exist_ok=True)
        print(f"Created save directory at {args.save_path}")

        with open(os.path.join(args.save_path, "training_args.json"), "w") as f:
            json.dump(vars(args), f, indent=4)

        log_file = os.path.join(args.save_path, f"training_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s [%(levelname)s] %(message)s",
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        logger = logging.getLogger(__name__)
        logger.info(f"Training arguments: {args}")
        print(f"Initialized distributed training with {world_size} GPUs")

    tokenizer = AutoTokenizer.from_pretrained(args.bert_model_id)

    if is_main_process:
        print("Loading datasets...")
    
    cdfg_list = torch.load("data/graph_list.pt", weights_only=False)
    with open("data/code_desc_list.json", "r") as f:
        code_and_des_list = json.load(f)

    full_dataset = MultiModalDataset(
        text_data=code_and_des_list,
        graph_data=cdfg_list,
        tokenizer=tokenizer,
        max_text_len=args.max_text_len,
        max_graph_len=args.max_graph_len
    )

    val_size = int(len(full_dataset) * args.val_split_ratio)
    train_size = len(full_dataset) - val_size
    train_dataset, val_dataset = random_split(
        full_dataset, 
        [train_size, val_size],
        generator=torch.Generator().manual_seed(42)
    )
    
    if is_main_process:
        logger.info(f"Dataset split: Train={len(train_dataset)}, Val={len(val_dataset)}")
    
    model_type = args.model_type
    if is_main_process:
        print("Creating models...")
    config = BertConfig.from_pretrained(args.bert_model_id) if model_type == "bert" else RobertaConfig.from_pretrained(args.bert_model_id) 

    original_max_length = tokenizer.model_max_length

    if args.max_text_len > original_max_length:
        config.max_position_embeddings = args.max_text_len if model_type == "bert" else args.max_text_len + 2
        bert_model = BertForPreTrainingWithMlmAndNspWithGraph.from_pretrained(args.bert_model_id, config=config, ignore_mismatched_sizes=True) \
            if model_type == "bert" else RobertaForMaskedLMWithGraph.from_pretrained(args.bert_model_id, config=config, ignore_mismatched_sizes=True)
        bert_model_aux = BertForPreTrainingWithMlmAndNspWithGraph.from_pretrained(args.bert_model_id) \
            if model_type == "bert" else RobertaForMaskedLMWithGraph.from_pretrained(args.bert_model_id)
        tokenizer.model_max_length = args.max_text_len
        original_embeddings = bert_model_aux.bert.embeddings.position_embeddings if model_type == "bert" else bert_model_aux.roberta.embeddings.position_embeddings
        if model_type == "bert":
            bert_model.bert.embeddings.position_embeddings = nn.Embedding(args.max_text_len, config.hidden_size)
            bert_model.bert.embeddings.position_embeddings.weight.data[:original_max_length, :] = original_embeddings.weight.data
        else:
            bert_model.roberta.embeddings.position_embeddings = nn.Embedding(args.max_text_len + 2, config.hidden_size)
            bert_model.roberta.embeddings.position_embeddings.weight.data[:original_max_length + 2, :] = original_embeddings.weight.data
        del bert_model_aux
    else:
        bert_model = BertForPreTrainingWithMlmAndNspWithGraph.from_pretrained(args.bert_model_id) \
            if model_type == "bert" else RobertaForMaskedLMWithGraph.from_pretrained(args.bert_model_id)
        
    cdfgnn = load_cdfgnn_model(args.cdfgnn_model_path)
    unirtl_model = UniRTL(bert_model=bert_model, cdfgnn=cdfgnn, num_node_classes=args.num_node_classes, graph_mask_prob=args.graph_mask_prob, \
        graph_loss_type=args.graph_loss_type, class_counts=compute_class_counts(train_dataset, args.num_node_classes)).to(device)
    cdfgnn_config = json.load(open(os.path.join(args.cdfgnn_model_path, "args.json")))
    unirtl_config = UniRTLConfig(bert_config=config, cdfgnn_config=cdfgnn_config, num_node_classes=args.num_node_classes, graph_mask_prob=args.graph_mask_prob)

    unirtl_model = torch.nn.parallel.DistributedDataParallel(
        unirtl_model, 
        device_ids=[rank],
        output_device=rank,
        find_unused_parameters=True
    )

    if is_main_process:
        print("Creating dataloaders...")

    train_sampler = DistributedSampler(
        train_dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True,
        drop_last=False
    )
    
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=args.per_device_batch_size,
        sampler=train_sampler,
        collate_fn=collate_fn,
        # num_workers=4,
        num_workers=0,
        # pin_memory=True,
        # persistent_workers=True
    )

    val_sampler = DistributedSampler(
        val_dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=False,
        drop_last=False
    )
    
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=args.per_device_batch_size,
        sampler=val_sampler,
        collate_fn=collate_fn,
        # num_workers=4,
        num_workers=0,
        # pin_memory=True,
        # persistent_workers=True
    )

    optimizer = torch.optim.AdamW(
        unirtl_model.parameters(), 
        lr=args.learning_rate,
        weight_decay=args.weight_decay
    )

    total_steps = len(train_dataloader) // args.gradient_accumulation_steps * args.epochs

    if args.lr_scheduler_type == "linear":
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=int(0.001 * total_steps),
            num_training_steps=total_steps
        )
    elif args.lr_scheduler_type == "cosine":
        scheduler = get_cosine_schedule_with_warmup(
            optimizer,
            num_warmup_steps=int(0.001 * total_steps),
            num_training_steps=total_steps
        )
    else:
        scheduler = None

    best_val_loss = float('inf')
    best_val_accuracy = 0.0

    for epoch in range(args.epochs):
        train_sampler.set_epoch(epoch)
        unirtl_model.train()
        total_train_loss = 0
        total_train_accuracy = 0
        step_train_losses = []
        step_train_accuracies = []
        
        if is_main_process:
            logger.info(f"\nStarting Training Epoch {epoch+1}/{args.epochs}")
        
        optimizer.zero_grad()

        accumulation_steps = 0
        
        for batch_idx, batch in enumerate(train_dataloader):
            inputs = convert_batch_to_unirtl_input_format(batch, tokenizer, mask_ratio=args.text_mask_ratio)
            inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v 
                     for k, v in inputs.items()}
            if model_type == "roberta":
                del inputs["next_sentence_label"]

            if 'graphs' in inputs:
                inputs['graphs'] = [graph.to(device) for graph in inputs['graphs']]

            outputs = unirtl_model(inputs)
            loss = outputs.loss / args.gradient_accumulation_steps 

            loss.backward()

            step_loss = loss.item() * args.gradient_accumulation_steps 
            step_accuracy = outputs.graph_accuracy.item() if outputs.graph_accuracy is not None else 0.0
            
            total_train_loss += step_loss
            total_train_accuracy += step_accuracy
            step_train_losses.append(step_loss)
            step_train_accuracies.append(step_accuracy)

            accumulation_steps += 1
            if accumulation_steps % args.gradient_accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(unirtl_model.parameters(), max_norm=args.grad_clip)

                optimizer.step()
                if scheduler is not None:
                    scheduler.step()
                optimizer.zero_grad()

            if is_main_process and (batch_idx % args.log_steps == 0 or batch_idx == len(train_dataloader)-1):
                avg_train_loss = total_train_loss / (batch_idx + 1)
                avg_train_accuracy = total_train_accuracy / (batch_idx + 1)
                current_lr = optimizer.param_groups[0]['lr']
                logger.info(f"Main Process(GPU {rank}) | Epoch {epoch+1}/{args.epochs} | Train Step {batch_idx}/{len(train_dataloader)} | "
                            f"Step Loss: {step_loss:.4f} | Step Text Loss: {outputs.text_loss.item():.4f} | Step Graph Loss: {outputs.graph_loss.item():.4f} | "
                            f"Step Graph Acc: {step_accuracy:.4f} | Avg Loss: {avg_train_loss:.4f} | Avg Acc: {avg_train_accuracy:.4f} | LR: {current_lr:.2e}")

        if accumulation_steps % args.gradient_accumulation_steps != 0:
            torch.nn.utils.clip_grad_norm_(unirtl_model.parameters(), max_norm=args.grad_clip)
            optimizer.step()
            if scheduler is not None:
                scheduler.step()
            optimizer.zero_grad()

        dist.barrier()
        avg_train_loss = total_train_loss / len(train_dataloader)
        avg_train_accuracy = total_train_accuracy / len(train_dataloader)
        
        total_train_loss_tensor = torch.tensor(avg_train_loss).to(device)
        total_train_accuracy_tensor = torch.tensor(avg_train_accuracy).to(device)
        
        dist.all_reduce(total_train_loss_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(total_train_accuracy_tensor, op=dist.ReduceOp.SUM)
        
        global_avg_train_loss = total_train_loss_tensor.item() / world_size
        global_avg_train_accuracy = total_train_accuracy_tensor.item() / world_size

        unirtl_model.eval()
        total_val_loss = 0
        total_val_accuracy = 0
        step_val_losses = []
        step_val_accuracies = []
        
        if is_main_process:
            logger.info(f"Starting Validation Epoch {epoch+1}/{args.epochs}")
        
        with torch.no_grad():
            for batch_idx, batch in enumerate(val_dataloader):
                inputs = convert_batch_to_unirtl_input_format(batch, tokenizer, mask_ratio=args.text_mask_ratio)
                inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v 
                         for k, v in inputs.items()}
                if model_type == "roberta":
                    del inputs["next_sentence_label"]

                if 'graphs' in inputs:
                    inputs['graphs'] = [graph.to(device) for graph in inputs['graphs']]

                outputs = unirtl_model(inputs)
                loss = outputs.loss

                step_loss = loss.item()
                step_accuracy = outputs.graph_accuracy.item() if outputs.graph_accuracy is not None else 0.0
                total_val_loss += step_loss
                total_val_accuracy += step_accuracy
                step_val_losses.append(step_loss)
                step_val_accuracies.append(step_accuracy)

                if is_main_process and (batch_idx % (max(1, args.log_steps // 5)) == 0 or batch_idx == len(val_dataloader)-1):
                    avg_val_loss = total_val_loss / (batch_idx + 1)
                    avg_val_accuracy = total_val_accuracy / (batch_idx + 1)
                    logger.info(f"Main Process(GPU {rank}) | Epoch {epoch+1}/{args.epochs} | Val Step {batch_idx}/{len(val_dataloader)} | "
                               f"Step Loss: {step_loss:.4f} | Step Text Loss: {outputs.text_loss.item():.4f} | Step Graph Loss: {outputs.graph_loss.item():.4f} | "
                               f"Step Graph Acc: {step_accuracy:.4f} | Avg Loss: {avg_val_loss:.4f} | Avg Acc: {avg_val_accuracy:.4f}")

        dist.barrier()
        avg_val_loss = total_val_loss / len(val_dataloader)
        avg_val_accuracy = total_val_accuracy / len(val_dataloader)
        
        total_val_loss_tensor = torch.tensor(avg_val_loss).to(device)
        total_val_accuracy_tensor = torch.tensor(avg_val_accuracy).to(device)
        
        dist.all_reduce(total_val_loss_tensor, op=dist.ReduceOp.SUM)
        dist.all_reduce(total_val_accuracy_tensor, op=dist.ReduceOp.SUM)
        
        global_avg_val_loss = total_val_loss_tensor.item() / world_size
        global_avg_val_accuracy = total_val_accuracy_tensor.item() / world_size
        
        if is_main_process:
            logger.info(f"Epoch {epoch+1} Completed | Train Loss: {global_avg_train_loss:.4f} | Val Loss: {global_avg_val_loss:.4f} | "
                       f"Train Acc: {global_avg_train_accuracy:.4f} | Val Acc: {global_avg_val_accuracy:.4f}")

            loss_log = {
                'epoch': epoch,
                'avg_train_loss': global_avg_train_loss,
                'avg_val_loss': global_avg_val_loss,
                'avg_train_accuracy': global_avg_train_accuracy,
                'avg_val_accuracy': global_avg_val_accuracy,
                'step_train_losses': step_train_losses,
                'step_val_losses': step_val_losses,
                'step_train_accuracies': step_train_accuracies,
                'step_val_accuracies': step_val_accuracies
            }
            with open(os.path.join(args.save_path, f"loss_epoch_{epoch+1}.json"), "w") as f:
                json.dump(loss_log, f, indent=4)

            if global_avg_val_loss < best_val_loss or global_avg_val_accuracy > best_val_accuracy:
                if global_avg_val_loss < best_val_loss:
                    best_val_loss = global_avg_val_loss
                if global_avg_val_accuracy > best_val_accuracy:
                    best_val_accuracy = global_avg_val_accuracy
                    
                best_model_path = os.path.join(args.save_path, "best")
                if not os.path.exists(best_model_path):
                    os.makedirs(best_model_path)
                save_model(unirtl_model.module, unirtl_config, tokenizer, best_model_path)
                logger.info(f"Saved best model with val loss {best_val_loss:.4f} and val accuracy {best_val_accuracy:.4f} to {best_model_path}")
    
    if is_main_process:
        final_model_path = os.path.join(args.save_path, "final")
        if not os.path.exists(final_model_path):
            os.makedirs(final_model_path)
        save_model(unirtl_model.module, unirtl_config, tokenizer, final_model_path)
        logger.info(f"Saved final model to {final_model_path}")
        logger.info("Training completed!")
    
    dist.destroy_process_group()

def load_cdfgnn_model(cdfgnn_model_path):
    args_dict = json.load(open(os.path.join(cdfgnn_model_path, "args.json")))
    args = argparse.Namespace(**args_dict)
    return load_model(args, cdfgnn_model_path)

def save_model(model, config, tokenizer, save_path):
    model.save_pretrained(save_path)
    config.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)

if __name__ == "__main__":
    main()