Merge pull request #9 from foundation-model-stack/validation_criteria

Robust level 1 validation testing

Author: JRosenkranz

aiu_fms_testing_utils/testing/validation.py

@@ -286,14 +286,15 @@ def capture_level_1_metrics(reference_logits_per_sentence, test_logits_per_sente
 
     return loss_metrics
 
-def filter_failed_level_1_cases(level_1_loss_metrics, fail_f):
+def filter_failed_level_1_cases(level_1_loss_metrics, fail_f, print_failed=False):
     failed_cases = []
     for (sentence_idx, token_idx, metrics_value) in level_1_loss_metrics:
         if fail_f(metrics_value):
             failed_cases.append((sentence_idx, token_idx, metrics_value))
-            print(
-                f"In sentence {sentence_idx+1}, the metric for token {token_idx} is {metrics_value}"
-            )
+            if print_failed:
+                dprint(
+                    f"In sentence {sentence_idx+1}, the metric for token {token_idx} is {metrics_value}"
+                )
     return failed_cases
 
 
@@ -304,4 +305,4 @@ def print_failed_cases(failed_cases, aiu_tokens, validation_tokens, tokenizer):
 
         aiu_str = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(aiu_token))
         validation_str = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(validation_token))
-        print(f"In sentence {sentence_index+1}/{len(aiu_tokens)}, token {token_index}, AIU outputs {aiu_token} instead of {validation_token} -- AIU val={aiu_str} != CPU val={validation_str}")
+        print(f"In sentence {sentence_index+1}/{len(aiu_tokens)}, token {token_index}, AIU outputs {aiu_token} instead of {validation_token} -- AIU val={aiu_str} -- CPU val={validation_str}")

aiu_fms_testing_utils/utils/__init__.py

@@ -1,8 +1,14 @@
 import torch
 import torch.nn as nn
 import time
+from fms.utils.tokenizers import BaseTokenizer
 from fms.utils.generation import generate
 from aiu_fms_testing_utils.utils.aiu_setup import dprint
+from typing import Optional, List, Tuple
+import os
+import requests
+import json
+import random
 
 def warmup_model(model: nn.Module, input_ids: torch.Tensor, max_new_tokens: int, **padding_kwargs):
     from torch_sendnn import torch_sendnn
@@ -25,4 +31,60 @@ def ids_for_prompt(prompt, tokenizer):
     if tokenizer.bos_token_id != tokenizer.eos_token_id:
         ids = [tokenizer.bos_token_id] + ids
     ids = torch.tensor(ids, dtype=torch.long, device="cpu")
-    return ids
+    return ids
+
+def __download_file(url, filename):
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        
+        with open(filename, 'wb') as file:
+            for chunk in response.iter_content(chunk_size=8192):
+                file.write(chunk)
+        print(f"Successfully downloaded {filename}")
+    
+    except requests.exceptions.RequestException as e:
+        print(f"An error occurred: {e}")
+
+def sample_sharegpt_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: BaseTokenizer,
+    prompt_length_min: int = 32,
+    prompt_length_max: int = 64,
+    seed: Optional[int] = None
+) -> List[Tuple[str, int]]:
+    if not os.path.exists(dataset_path):
+        print("downloading share-gpt dataset as it does not exist")
+        __download_file("https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json", dataset_path)
+
+    # Load the dataset.
+    with open(dataset_path, encoding='utf-8') as f:
+        dataset = json.load(f)
+    # Filter out the conversations with less than 2 turns.
+    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
+    # Only keep the first two turns of each conversation.
+    dataset = [(data["conversations"][0]["value"],
+                data["conversations"][1]["value"]) for data in dataset]
+
+    # Shuffle the dataset.
+    if seed is not None:
+        random.Random(seed).shuffle(dataset)
+
+    # Filter out sequences that are too long or too short
+    filtered_dataset: List[Tuple[str, int, int]] = []
+    for i in range(len(dataset)):
+        if len(filtered_dataset) == num_requests:
+            break
+
+        # Tokenize the prompts and completions.
+        prompt = dataset[i][0]
+        prompt_token_ids = ids_for_prompt(prompt, tokenizer)
+        
+        prompt_len = len(prompt_token_ids)
+        if prompt_len < prompt_length_min or prompt_len > prompt_length_max:
+            # Prune too short or too long sequences.
+            continue
+        filtered_dataset.append((prompt, prompt_len))
+    
+    return filtered_dataset

scripts/generate_metrics.py

@@ -1,4 +1,5 @@
 import argparse
+import ast
 import json
 import os
 import random
@@ -8,7 +9,7 @@
 
 from aiu_fms_testing_utils.testing.validation import capture_level_1_metrics, extract_validation_information, LogitsExtractorHook, print_failed_cases, \
     validate_level_0, GoldenTokenHook, top_k_loss_calculator
-from aiu_fms_testing_utils.utils import ids_for_prompt
+from aiu_fms_testing_utils.utils import ids_for_prompt, sample_sharegpt_requests
 from fms.models import get_model
 from fms.utils import tokenizers
 from fms.utils.generation import pad_input_ids
@@ -84,10 +85,30 @@
     help="top k values per token to generate loss on",
     default=20
 )
+parser.add_argument(
+    "--num_test_tokens_per_sequence",
+    type=int,
+    help="number of tokens in test. For instance, if max_new_tokens=128 and num_test_tokens_per_sequence=256, this means we will generate data over 2 sample prompts. If not set, will be set to max_new_tokens",
+    default=None
+)
+parser.add_argument(
+    "--extra_get_model_kwargs",
+    nargs='*',
+    default={},
+    help="Use this to override model configuration values to get model. Example: --extra_get_model_kwargs nlayers=2,..."
+)
 args = parser.parse_args()
 
+extra_get_model_kwargs = {}
+for a in args.extra_get_model_kwargs:
+     a_split = a.split("=")
+     try:
+        extra_get_model_kwargs[a_split[0]] = ast.literal_eval(a_split[1])
+     except ValueError:
+        extra_get_model_kwargs[a_split[0]] = a_split[1]
 
-prefix = f"{args.variant.replace('/', '--')}_max-new-tokens-{args.max_new_tokens}_batch-size-{args.batch_size}_seq-length{args.min_pad_length}_dtype-{args.default_dtype}"
+# this follows the same pattern of naming in test_shapes. This way we can save and re-use for quicker shape testing.
+prefix = f"{args.variant.replace('/', '--')}_max-new-tokens-{args.max_new_tokens}_batch-size-{args.batch_size}_seq-length-{args.min_pad_length}_dtype-{args.default_dtype}"
 if os.path.exists(os.path.join(args.output_dir, f"{prefix}.prob_mean.csv")):
     print("skipping metric generation as it has already been done")
     exit(0)
@@ -115,11 +136,11 @@
     model_path=args.model_path,
     device_type="cuda",
     data_type=default_dtype,
+    **extra_get_model_kwargs,
 )
 
-print("loaded cuda model")
-
 cuda_model.eval()
+print("loaded cuda model")
 
 # prepare the cpu model (this is the reference)
 cpu_model = get_model(
@@ -128,45 +149,11 @@
     model_path=args.model_path,
     device_type="cpu",
     data_type=torch.float32,
+    **extra_get_model_kwargs,
 )
 cpu_model.eval()
 print("loaded cpu model")
 
-def sample_sharegpt_requests(
-    dataset_path: str,
-    num_requests: int,
-    tokenizer,
-) -> List[Tuple[str, int, int, None]]:
-    # Load the dataset.
-    with open(dataset_path, encoding='utf-8') as f:
-        dataset = json.load(f)
-    # Filter out the conversations with less than 2 turns.
-    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
-    # Only keep the first two turns of each conversation.
-    dataset = [(data["conversations"][0]["value"],
-                data["conversations"][1]["value"]) for data in dataset]
-
-    # Shuffle the dataset.
-    random.Random(42).shuffle(dataset)
-
-    # Filter out sequences that are too long or too short
-    filtered_dataset: List[Tuple[str, int, int]] = []
-    for i in range(len(dataset)):
-        if len(filtered_dataset) == num_requests:
-            break
-
-        # Tokenize the prompts and completions.
-        prompt = dataset[i][0]
-        prompt_token_ids = ids_for_prompt(prompt, tokenizer)
-        
-        prompt_len = len(prompt_token_ids)
-        if prompt_len < 32 or prompt_len > args.min_pad_length:
-            # Prune too short sequences.
-            continue
-        filtered_dataset.append((prompt, prompt_len))
-
-    return filtered_dataset
-
 def find_eos_index(reference_tokens, eos_token_id):
     result = []
     for sentence in reference_tokens:
@@ -184,21 +171,17 @@ def filter_before_eos(l, filter_indexes):
     from itertools import groupby
     filtered_results = [list(g)[:filter_indexes[k]] for k, g in groupby(l, key=lambda x: x[0])]
     return [item for sublist in filtered_results for item in sublist]
-            
-prompts_and_lens = sample_sharegpt_requests(args.sharegpt_path, args.batch_size, tokenizer)
-print(f"prompt_lengths: {[pl[1] for pl in prompts_and_lens]}")
-prompts = [ids_for_prompt(pl[0], tokenizer) for pl in prompts_and_lens]
 
-padding_length = args.min_pad_length
+def __prepare_inputs(batch_size, seq_length, tokenizer, seed=0):
+    prompts_and_sizes = sample_sharegpt_requests(args.sharegpt_path, batch_size, tokenizer, seq_length // 2, seq_length, seed)
+    prompt_list = []
+    for prompt, _ in prompts_and_sizes:
+        prompt_list.append(ids_for_prompt(prompt, tokenizer))
 
-has_padding = args.batch_size > 1 or padding_length != 0
-max_len = max([len(prompt) for prompt in prompts])
+    input_ids, padding_kwargs = pad_input_ids(prompt_list, min_pad_length=seq_length)
+    return input_ids, padding_kwargs
 
-if has_padding:
-    ids, padding_kwargs = pad_input_ids(prompts, min_pad_length=padding_length)
-else:
-    ids = prompts
-    padding_kwargs = {}
+ids, padding_kwargs = __prepare_inputs(args.batch_size, args.min_pad_length, tokenizer)
 
 # first test validation level 0
 cpu_validation_info = extract_validation_information(
@@ -231,63 +214,88 @@ def filter_before_eos(l, filter_indexes):
 if len(failed_responses) != 0:    
     print_failed_cases(failed_responses, cpu_static_tokens, cuda_static_tokens, tokenizer)
 
-# generate aiu validation info
-cuda_validation_info = extract_validation_information(
-    cuda_model,
-    ids.to("cuda"),
-    args.max_new_tokens,
-    GoldenTokenHook(cpu_static_tokens, "cuda"),
-    only_last_token=True,
-    **{k: v.to("cuda") for k,v in padding_kwargs.items()}
-)
-
-print("extracted cuda validation information level 1")
-
-cross_entropy = lambda r, t: torch.nn.CrossEntropyLoss()(r, t.softmax(dim=1).to(dtype=torch.float32))
-prob_mean = lambda r, t: torch.mean((r.softmax(dim=1).to(dtype=torch.float32) / t.softmax(dim=1).to(dtype=torch.float32)) - 1.0)
-prob_std = lambda r, t: torch.std(r.softmax(dim=1).to(dtype=torch.float32) / t.softmax(dim=1).to(dtype=torch.float32))
-diff_mean = lambda r, t: torch.mean(r.softmax(dim=1).to(dtype=torch.float32) - t.softmax(dim=1).to(dtype=torch.float32))
-
 def write_csv(l, path, metric):
     with open(path, 'w') as f:
         f.write(f'{metric}\n')
         for t in l:
             f.write(f"{t[2].item()}\n") 
         f.close()
 
-prefix = f"{args.variant.replace('/', '--')}_max-new-tokens-{args.max_new_tokens}_batch-size-{args.batch_size}_seq-length{args.min_pad_length}_dtype-{args.default_dtype}"
-
-cpu_validation_info.save(os.path.join(args.output_dir, f"{prefix}.cpu_output_logits.out"))
-cuda_validation_info.save(os.path.join(args.output_dir, f"{prefix}.cuda_output_logits.out"))
+num_test_tokens_per_sequence = args.num_test_tokens_per_sequence
+if num_test_tokens_per_sequence is None:
+    num_test_tokens_per_sequence = args.max_new_tokens
 
-level_1_metrics = capture_level_1_metrics(
-    cpu_validation_info.get_info("logits"),
-    cuda_validation_info.get_info("logits"),
-    top_k_loss_calculator(args.topk_per_token, prob_mean),
-)
-loss_metrics = filter_before_eos(level_1_metrics, eos_indexes)
-write_csv(loss_metrics, os.path.join(args.output_dir, f"{prefix}.prob_mean.csv"), "prob_mean")
+cross_entropy = lambda r, t: torch.nn.CrossEntropyLoss()(r, t.softmax(dim=1).to(dtype=torch.float32))
+prob_mean = lambda r, t: torch.mean((r.softmax(dim=1).to(dtype=torch.float32) / t.softmax(dim=1).to(dtype=torch.float32)) - 1.0)
+prob_std = lambda r, t: torch.std(r.softmax(dim=1).to(dtype=torch.float32) / t.softmax(dim=1).to(dtype=torch.float32))
+diff_mean = lambda r, t: torch.mean(r.softmax(dim=1).to(dtype=torch.float32) - t.softmax(dim=1).to(dtype=torch.float32))
 
-level_1_metrics = capture_level_1_metrics(
-    cpu_validation_info.get_info("logits"),
-    cuda_validation_info.get_info("logits"),
-    top_k_loss_calculator(args.topk_per_token, prob_std),
-)
-loss_metrics = filter_before_eos(level_1_metrics, eos_indexes)
-write_csv(loss_metrics, os.path.join(args.output_dir, f"{prefix}.prob_std.csv"), "prob_std")
+prob_mean_metrics = []
+prob_std_metrics = []
+prob_diff_metrics = []
+prob_ce_loss_metrics = []
 
-level_1_metrics = capture_level_1_metrics(
-    cpu_validation_info.get_info("logits"),
-    cuda_validation_info.get_info("logits"),
-    top_k_loss_calculator(args.topk_per_token, cross_entropy),
-)
-loss_metrics = filter_before_eos(level_1_metrics, eos_indexes)
-write_csv(loss_metrics, os.path.join(args.output_dir, f"{prefix}.ce.csv"), "ce")
+prefix = f"{args.variant.replace('/', '--')}_max-new-tokens-{args.max_new_tokens}_batch-size-{args.batch_size}_seq-length{args.min_pad_length}_dtype-{args.default_dtype}"
 
-level_1_metrics = capture_level_1_metrics(
-    cpu_validation_info.get_info("logits"),
-    cuda_validation_info.get_info("logits"),
-    top_k_loss_calculator(args.topk_per_token, diff_mean),
-)
-loss_metrics = filter_before_eos(level_1_metrics, eos_indexes)
-write_csv(loss_metrics, os.path.join(args.output_dir, f"{prefix}.diff_mean.csv"), "diff_mean")
+for i in range(num_test_tokens_per_sequence // args.max_new_tokens):
+    ids, padding_kwargs = __prepare_inputs(args.batch_size, args.min_pad_length, tokenizer, i)
+
+    # only need to compute this once if we aren't generating more test data
+    if num_test_tokens_per_sequence > args.max_new_tokens:
+        cpu_validation_info = extract_validation_information(
+            cpu_model,
+            ids,
+            args.max_new_tokens,
+            LogitsExtractorHook(),
+            attn_algorithm="math",
+            **padding_kwargs
+        )
+        eos_indexes = find_eos_index(cpu_validation_info.get_info("tokens"), tokenizer.eos_token_id)
+
+    # generate aiu validation info
+    cuda_validation_info = extract_validation_information(
+        cuda_model,
+        ids.to("cuda"),
+        args.max_new_tokens,
+        GoldenTokenHook(cpu_validation_info.get_info("tokens"), "cuda"),
+        only_last_token=True,
+        **{k: v.to("cuda") for k,v in padding_kwargs.items()}
+    )
+
+    print("extracted cuda validation information level 1")
+
+    cpu_validation_info.save(os.path.join(args.output_dir, f"{prefix}.cpu_validation_info.{i}.out"))
+    cuda_validation_info.save(os.path.join(args.output_dir, f"{prefix}.cuda_validation_info.{i}.out"))
+
+    level_1_metrics = capture_level_1_metrics(
+        cpu_validation_info.get_info("logits"),
+        cuda_validation_info.get_info("logits"),
+        top_k_loss_calculator(args.topk_per_token, prob_mean),
+    )
+    prob_mean_metrics.extend(filter_before_eos(level_1_metrics, eos_indexes))
+
+    level_1_metrics = capture_level_1_metrics(
+        cpu_validation_info.get_info("logits"),
+        cuda_validation_info.get_info("logits"),
+        top_k_loss_calculator(args.topk_per_token, prob_std),
+    )
+    prob_std_metrics.extend(filter_before_eos(level_1_metrics, eos_indexes))
+
+    level_1_metrics = capture_level_1_metrics(
+        cpu_validation_info.get_info("logits"),
+        cuda_validation_info.get_info("logits"),
+        top_k_loss_calculator(args.topk_per_token, cross_entropy),
+    )
+    prob_ce_loss_metrics.extend(filter_before_eos(level_1_metrics, eos_indexes))
+
+    level_1_metrics = capture_level_1_metrics(
+        cpu_validation_info.get_info("logits"),
+        cuda_validation_info.get_info("logits"),
+        top_k_loss_calculator(args.topk_per_token, diff_mean),
+    )
+    prob_diff_metrics.extend(filter_before_eos(level_1_metrics, eos_indexes))
+
+write_csv(prob_mean_metrics, os.path.join(args.output_dir, f"{prefix}.prob_mean.csv"), "prob_mean")
+write_csv(prob_std_metrics, os.path.join(args.output_dir, f"{prefix}.prob_std.csv"), "prob_std")
+write_csv(prob_ce_loss_metrics, os.path.join(args.output_dir, f"{prefix}.ce.csv"), "ce")
+write_csv(prob_diff_metrics, os.path.join(args.output_dir, f"{prefix}.diff_mean.csv"), "diff_mean")