diff --git a/logger.py b/logger.py index 13e465ec578b1e3ce02a3cf154892c211a03bce0..794367c7497c14f9ca91295c46732c931bf24315 100644 --- a/logger.py +++ b/logger.py @@ -2,8 +2,11 @@ import torch import psutil import csv import math +import datetime + +NOW_STRING = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") +CSV_FILE = f"output_{NOW_STRING}.csv" -CSV_FILE = "output.csv" def init_csv(): """Initialize CSV file with headers.""" @@ -13,6 +16,7 @@ def init_csv(): "peak_memory_usage_allocated_GB", "peak_memory_usage_reserved_GB", "loss", "perplexity"]) + def measure_memory(): """Measure memory usage from CUDA or CPU.""" if torch.cuda.is_available(): @@ -28,6 +32,7 @@ def measure_memory(): return peak_history, max_allocated, max_reserved + def log_to_csv(epoch, step, compute_time, loss): """Log training metrics to CSV file.""" peak_history, max_allocated, max_reserved = measure_memory() diff --git a/main.py b/main.py index 2cc320355402aab72809d73cbe631a2f69d9cc69..b0773489cfbac6b38bb37ae628aa278c23f48b6f 100644 --- a/main.py +++ b/main.py @@ -72,13 +72,6 @@ if __name__ == "__main__": dataset = load_data(args, tokenizer) - optimizer, model = get_optimizer(args, model) - - num_steps = ceil(args.num_epochs * len(dataset) / args.batch_size) - scheduler = get_scheduler( - optimizer, args.lr_scheduler, args.warm_up_fraction, num_steps, args.lr, args.lr_min - ) - shuffle = True if args.shuffle == "true" else False if args.mode == "pretraining": dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=shuffle) @@ -87,6 +80,13 @@ if __name__ == "__main__": else: raise ValueError("Invalid mode. Choose 'pretraining' or 'finetuning'") + optimizer, model = get_optimizer(args, model) + + num_steps = ceil(args.num_epochs * len(dataloader)) + scheduler = get_scheduler( + optimizer, args.lr_scheduler, args.warm_up_fraction, num_steps, args.lr, args.lr_min + ) + trained_model = train(device, accelerator, scheduler, model, optimizer, dataloader, num_epochs=args.num_epochs) file_name = f"{args.model}_{args.optimizer}_pretrained" if args.mode == "pretraining" else f"{args.model}_{args.optimizer}_finetuned" diff --git a/requirements.txt b/requirements.txt index 022097615bd165545cfa6c1325d860ec3aaa8ae1..3ca4858725aeb28e2fdc168ec1b54d6df57529fe 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,8 @@ -galore-torch -datasets -transformers -torch -accelerate -psutil -peft -argparse \ No newline at end of file +galore-torch==1.0 +datasets==3.1.0 +transformers==4.46.3 +torch==2.4.1 +accelerate==1.0.1 +psutil==7.0.0 +peft==0.13.2 +argparse==1.4.0 \ No newline at end of file