GRAG-main/train_local.py

import os

# 设置可见的GPU设备 - 新增这行
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  # 只使用设备0和1

# 设置Hugging Face国内镜像源 (必须在import transformers之前设置)
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['HF_HUB_ENABLE_HF_TRANSFER'] = '1'

# 强制离线模式 - 确保从本地加载模型而不联网
os.environ['TRANSFORMERS_OFFLINE'] = '1'
os.environ['HF_DATASETS_OFFLINE'] = '1'

# 新增：设置本地WandB服务地址（可选但推荐）
os.environ['WANDB_BASE_URL'] = 'http://117.72.35.222:8081'  # 添加这一行
# 设置WandB API密钥 (替换为您的实际密钥)
os.environ['WANDB_API_KEY'] = 'local-84f1b57b7ab8df20a8881450fcfe22acb87bf9e2'  # 替换为您的本地WandB API密钥

import wandb
import gc
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from torch.nn.utils import clip_grad_norm_

from src.model import load_model, llama_model_path
from src.dataset import load_dataset
from src.utils.evaluate import eval_funcs
from src.utils.config import parse_args_llama
from src.utils.ckpt import _save_checkpoint, _reload_best_model
from src.utils.collate import collate_fn
from src.utils.seed import seed_everything
from src.utils.lr_schedule import adjust_learning_rate

# 初始化CUDA状态 - 新增这行
if torch.cuda.is_available():
    torch.cuda.init()  # 确保CUDA正确初始化



def main(args):
    # 打印可用设备信息 - 新增这行
    print(f"可用GPU设备数量: {torch.cuda.device_count()}")
    print(f"当前设备: {torch.cuda.current_device()}")

    # Step 1: Set up wandb
    seed = args.seed
    wandb.init(project=f"{args.project}",
               name=f"{args.dataset}_{args.model_name}_seed{seed}",
               config=args)

    seed_everything(seed=args.seed)
    print(args)

    dataset = load_dataset[args.dataset]()
    idx_split = dataset.get_idx_split()

    # Step 2: Build Dataset
    train_dataset = [dataset[i] for i in idx_split['train']]
    val_dataset = [dataset[i] for i in idx_split['val']]
    test_dataset = [dataset[i] for i in idx_split['test']]

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, drop_last=True, pin_memory=True, shuffle=True,
                              collate_fn=collate_fn)
    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, drop_last=False, pin_memory=True, shuffle=False,
                            collate_fn=collate_fn)
    test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size, drop_last=False, pin_memory=True,
                             shuffle=False, collate_fn=collate_fn)

    # 使用本地模型路径
    local_model_path = args.local_model_path
    print(f"使用本地模型路径: {local_model_path}")
    print(f"确保该路径包含完整的模型文件: config.json, pytorch_model.bin 等")

    # 设置模型路径
    args.llm_model_path = local_model_path

    # Step 3: Build Model
    # 添加显存优化选项 - 修改这行
    model = load_model[args.model_name](
        graph_type=dataset.graph_type,
        args=args,
        init_prompt=dataset.prompt,
        device_map="auto"  # 自动分配设备
    )

    # 打印模型所在设备 - 新增这行
    print(f"模型加载到设备: {model.device}")

    # Step 4 Set Optimizer
    params = [p for _, p in model.named_parameters() if p.requires_grad]
    optimizer = torch.optim.AdamW(
        [{'params': params, 'lr': args.lr, 'weight_decay': args.wd}, ],
        betas=(0.9, 0.95)
    )
    trainable_params, all_param = model.print_trainable_params()
    print(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}")

    # Step 5. Training
    num_training_steps = args.num_epochs * len(train_loader)
    progress_bar = tqdm(range(num_training_steps))
    best_val_loss = float('inf')
    #best_epoch = 0

    for epoch in range(args.num_epochs):
        # 清空显存缓存 - 新增这行
        #torch.cuda.empty_cache()

        model.train()
        epoch_loss, accum_loss = 0., 0.

        for step, batch in enumerate(train_loader):
            # 确保数据在正确设备上 - 新增这行
            if hasattr(model, 'device'):
                batch = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}

            optimizer.zero_grad()
            loss = model(batch)
            loss.backward()

            clip_grad_norm_(optimizer.param_groups[0]['params'], 0.1)

            if (step + 1) % args.grad_steps == 0:
                adjust_learning_rate(optimizer.param_groups[0], args.lr, step / len(train_loader) + epoch, args)

            optimizer.step()
            epoch_loss, accum_loss = epoch_loss + loss.item(), accum_loss + loss.item()

            if (step + 1) % args.grad_steps == 0:
                lr = optimizer.param_groups[0]["lr"]
                wandb.log({'Lr': lr})
                wandb.log({'Accum Loss': accum_loss / args.grad_steps})
                accum_loss = 0.

            progress_bar.update(1)

        print(f"Epoch: {epoch}|{args.num_epochs}: Train Loss (Epoch Mean): {epoch_loss / len(train_loader)}")
        wandb.log({'Train Loss (Epoch Mean)': epoch_loss / len(train_loader)})

        val_loss = 0.
        eval_output = []
        model.eval()
        with torch.no_grad():
            for step, batch in enumerate(val_loader):
                # 确保数据在正确设备上 - 新增这行
                if hasattr(model, 'device'):
                    batch = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}

                loss = model(batch)
                val_loss += loss.item()
            val_loss = val_loss / len(val_loader)
            print(f"Epoch: {epoch}|{args.num_epochs}: Val Loss: {val_loss}")
            wandb.log({'Val Loss': val_loss})

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            _save_checkpoint(model, optimizer, epoch, args, is_best=True)
            best_epoch = epoch

        print(f'Epoch {epoch} Val Loss {val_loss} Best Val Loss {best_val_loss} Best Epoch {best_epoch}')

        if epoch - best_epoch >= args.patience:
            print(f'Early stop at epoch {epoch}')
            break

    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()

    # Step 5. Evaluating
    model = _reload_best_model(model, args)
    model.eval()
    eval_output = []
    progress_bar_test = tqdm(range(len(test_loader)))
    for step, batch in enumerate(test_loader):
        with torch.no_grad():
            # 确保数据在正确设备上 - 新增这行
            if hasattr(model, 'device'):
                batch = {k: v.to(model.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}

            output = model.inference(batch)
            eval_output.append(output)

        progress_bar_test.update(1)

    # Step 6. Post-processing & compute metrics
    os.makedirs(f'{args.output_dir}/{args.dataset}', exist_ok=True)
    path = f'{args.output_dir}/{args.dataset}/model_name_{args.model_name}_llm_model_name_{args.llm_model_name}_llm_frozen_{args.llm_frozen}_max_txt_len_{args.max_txt_len}_max_new_tokens_{args.max_new_tokens}_gnn_model_name_{args.gnn_model_name}_patience_{args.patience}_num_epochs_{args.num_epochs}_seed{seed}.csv'
    acc = eval_funcs[args.dataset](eval_output, path)
    print(f'Test Acc {acc}')
    wandb.log({'Test Acc': acc})


if __name__ == "__main__":
    # 确保CUDA正确初始化 - 新增这行
    if torch.cuda.is_available():
        torch.cuda.init()

    args = parse_args_llama()

    # 添加本地模型路径参数
    if not hasattr(args, 'local_model_path'):
        args.local_model_path = "/data2/ycj/models/Llama-2-7b-chat-hf"  # 默认路径

    main(args)
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    gc.collect()