Spaces:
Sleeping
Sleeping
| import json | |
| import os | |
| import random | |
| import time | |
| import torch | |
| from distributed_training.data.dataset import DataLoader | |
| from huggingface_hub import create_tag, list_repo_refs, scan_cache_dir | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| device = "cuda" | |
| test_indices_length = 1000 | |
| AUTOMATE = True | |
| models = [ | |
| "distributed/optimized-gpt2-2b", | |
| "distributed/optimized-gpt2-1b", | |
| "distributed/optimized-gpt2-500m", | |
| "distributed/optimized-gpt2-250m", | |
| "distributed/optimized-gpt2-250m-v0.1.3", | |
| "distributed/optimized-gpt2-250m-v0.1.1", | |
| "distributed/gpt2-94m", | |
| ] | |
| if os.path.exists("results.json"): | |
| with open("results.json", "r") as file: | |
| results = json.load(file) | |
| else: | |
| results = {} | |
| while True: | |
| for model_name in [models[0]]: | |
| if model_name not in results.keys(): | |
| results[model_name] = {} | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| "distributed/optimized-gpt2-250m", trust_remote_code=True | |
| ) | |
| refs = list_repo_refs(model_name, repo_type="model") | |
| global_epoch = max([int(tag.name) for tag in refs.tags]) if refs.tags else None | |
| if global_epoch in results[model_name]["main-net"].keys(): | |
| print(f"Results for epoch {global_epoch} already calcualted") | |
| time.sleep(30 * 60) | |
| for epoch in range(0, global_epoch, 1): | |
| if str(epoch) in results[model_name]["main-net"].keys(): | |
| continue | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, revision=str(epoch), trust_remote_code=True | |
| ) | |
| model = model.to(device) | |
| search_start = random.choice( | |
| range(DataLoader.max_pages - test_indices_length + 1) | |
| ) | |
| group = [i for i in range(search_start, search_start + test_indices_length)] | |
| dataloader = DataLoader( | |
| batch_size=1, | |
| sequence_length=1024, | |
| rows=group, | |
| ) | |
| total_loss = 0 | |
| index = 0 | |
| # Train data for one epoch | |
| for index, batch in enumerate(dataloader): | |
| inputs = batch[0].to(device) | |
| labels = batch[1].to(device) | |
| if len(inputs[0]) != len(labels[0]): | |
| breakpoint() | |
| if "optimized" in model_name: | |
| outputs = model(input_ids=inputs, labels=labels) | |
| loss = outputs[1] | |
| else: | |
| outputs = model(input_ids=inputs, labels=inputs) | |
| loss = outputs.loss | |
| # Accumulate Total Loss | |
| total_loss += loss.detach().item() | |
| # Backward Pass | |
| model.zero_grad() | |
| average_loss = total_loss / (index + 1) | |
| results[model_name]["main-net"][str(epoch)] = [average_loss] | |
| print(f"Epoch: {epoch} Average Loss: {average_loss:.2f}") | |
| with open("results.json", "w") as outfile: | |
| json.dump(results, outfile, indent=4) | |
| current_revision = model.config._commit_hash | |
| keep_recent = 1 | |
| try: | |
| cache_info = scan_cache_dir() | |
| for repo in cache_info.repos: | |
| if repo.repo_id == model_name: | |
| revisions = sorted( | |
| repo.revisions, key=lambda r: r.last_modified, reverse=True | |
| ) | |
| current_index = next( | |
| ( | |
| i | |
| for i, r in enumerate(revisions) | |
| if r.commit_hash == current_revision | |
| ), | |
| None, | |
| ) | |
| if current_index is not None: | |
| for revision in revisions[ | |
| max(current_index + 1, keep_recent) : | |
| ]: | |
| cache_info.delete_revisions( | |
| revision.commit_hash | |
| ).execute() | |
| break | |
| except: | |
| print( | |
| "Failed to delete previous model version from cache. This might lead to 100% disk space utlisation in the future." | |
| ) | |