From 3862933d297225aec26344f930fe7fb57a29395b Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:18:36 +0530 Subject: [PATCH 01/13] visualize the backward graph of specific blocks (cherry picked from commit dc3a7e84d45ff2a220a6a3772279c795523dbb8b) --- experiments/util/visualize_graph.py | 51 +++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 experiments/util/visualize_graph.py diff --git a/experiments/util/visualize_graph.py b/experiments/util/visualize_graph.py new file mode 100644 index 0000000..0454f01 --- /dev/null +++ b/experiments/util/visualize_graph.py @@ -0,0 +1,51 @@ +# ruff: noqa +import argparse + +import torch +from torchview import draw_graph +from torchviz import make_dot + +from experiments.util import models +import torchvision.models as tvm +import transformers.models as tfm +from transformers import AutoConfig + +to_test = { + 'bert_encoder': lambda: models.transformer_model_fns['bert']().bert.encoder.layer[0], + 'memsave_bert_encoder': lambda: models.transformer_model_fns['memsave_bert']().bert.encoder.layer[0], + 'bart_encoder': lambda: models.transformer_model_fns['bart']().decoder.layers[0], + 'memsave_bart_encoder': lambda: models.transformer_model_fns['memsave_bart']().decoder.layers[0], + 'gpt2_layer': lambda: models.transformer_model_fns['gpt2']().transformer.h[0], + 'memsave_gpt2_layer': lambda: models.transformer_model_fns['memsave_gpt2']().transformer.h[0], + 't5_decoder': lambda: models.transformer_model_fns['t5']().decoder.block[1], + 'memsave_t5_decoder': lambda: models.transformer_model_fns['memsave_t5']().decoder.block[1], +} + +def run_single(model_fn, x): + model = model_fn() + + y = model(x) + dot = make_dot( + y.mean(), + params=dict(model.named_parameters()), + show_attrs=True, + show_saved=True, + ) + dot.render(filename=args.model, directory="torchviz-output") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="deeprelumodel", help="Which model to use" + ) + + args = parser.parse_args() + + models.conv_input_shape = (3, 64, 64) + x = torch.rand(7, *models.conv_input_shape) + + model = models.conv_model_fns.get(args.model) + assert model is not None From 9d94befdf9de51c9de2305cc5f1e998eb318bca4 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Tue, 7 May 2024 19:44:16 +0530 Subject: [PATCH 02/13] minor (cherry picked from commit c0fdf212513ff96a725e21effe3004092ed7ec74) --- experiments/util/visualize_graph.py | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/experiments/util/visualize_graph.py b/experiments/util/visualize_graph.py index 0454f01..a23e592 100644 --- a/experiments/util/visualize_graph.py +++ b/experiments/util/visualize_graph.py @@ -21,7 +21,7 @@ 'memsave_t5_decoder': lambda: models.transformer_model_fns['memsave_t5']().decoder.block[1], } -def run_single(model_fn, x): +def run_single(model_fn, name, x): model = model_fn() y = model(x) @@ -31,21 +31,24 @@ def run_single(model_fn, x): show_attrs=True, show_saved=True, ) - dot.render(filename=args.model, directory="torchviz-output") + dot.render(filename=name, directory="torchviz-output") if __name__ == "__main__": - import argparse + # import argparse - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", type=str, default="deeprelumodel", help="Which model to use" - ) + # parser = argparse.ArgumentParser() + # parser.add_argument( + # "--model", type=str, default="deeprelumodel", help="Which model to use" + # ) + + # args = parser.parse_args() - args = parser.parse_args() + # models.conv_input_shape = (3, 64, 64) + models.transformer_input_shape = (5000, 1024) - models.conv_input_shape = (3, 64, 64) - x = torch.rand(7, *models.conv_input_shape) + for name,model_fn in to_test.items(): + x = torch.rand(7, *models.transformer_input_shape) - model = models.conv_model_fns.get(args.model) - assert model is not None + run_single(model_fn, name, x) + \ No newline at end of file From b2ae54fef2e3a29723bdc23db65ab43c45d20f21 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Tue, 7 May 2024 20:28:12 +0530 Subject: [PATCH 03/13] fixes for vis minor (cherry picked from commit 4410b5f5615fc07d0b9aed5d3b2888f26f7ec8eb) --- experiments/util/visualize_graph.py | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/experiments/util/visualize_graph.py b/experiments/util/visualize_graph.py index a23e592..a74a27c 100644 --- a/experiments/util/visualize_graph.py +++ b/experiments/util/visualize_graph.py @@ -11,19 +11,17 @@ from transformers import AutoConfig to_test = { - 'bert_encoder': lambda: models.transformer_model_fns['bert']().bert.encoder.layer[0], - 'memsave_bert_encoder': lambda: models.transformer_model_fns['memsave_bert']().bert.encoder.layer[0], - 'bart_encoder': lambda: models.transformer_model_fns['bart']().decoder.layers[0], - 'memsave_bart_encoder': lambda: models.transformer_model_fns['memsave_bart']().decoder.layers[0], - 'gpt2_layer': lambda: models.transformer_model_fns['gpt2']().transformer.h[0], - 'memsave_gpt2_layer': lambda: models.transformer_model_fns['memsave_gpt2']().transformer.h[0], - 't5_decoder': lambda: models.transformer_model_fns['t5']().decoder.block[1], - 'memsave_t5_decoder': lambda: models.transformer_model_fns['memsave_t5']().decoder.block[1], + 'bert_encoder': ['bert', lambda model: model.bert.encoder.layer[0]], + 'memsave_bert_encoder': ['memsave_bert', lambda model: model.bert.encoder.layer[0]], + 'bart_encoder': ['bart', lambda model: model.decoder.layers[0]], + 'memsave_bart_encoder': ['memsave_bart', lambda model: model.decoder.layers[0]], + 'gpt2_layer': ['gpt2', lambda model: model.transformer.h[0]], + 'memsave_gpt2_layer': ['memsave_gpt2', lambda model: model.transformer.h[0]], + 't5_decoder': ['t5', lambda model: model.decoder.block[1]], + 'memsave_t5_decoder': ['memsave_t5', lambda model: model.decoder.block[1]], } -def run_single(model_fn, name, x): - model = model_fn() - +def run_single(model, name, x): y = model(x) dot = make_dot( y.mean(), @@ -47,8 +45,13 @@ def run_single(model_fn, name, x): # models.conv_input_shape = (3, 64, 64) models.transformer_input_shape = (5000, 1024) - for name,model_fn in to_test.items(): + for name in to_test: + model_name, block_fn = to_test[name] + config = models.get_transformers_config(model_name) + + models.transformer_input_shape = (config.vocab_size, config.hidden_size) x = torch.rand(7, *models.transformer_input_shape) - run_single(model_fn, name, x) + model = models.transformer_model_fns.get(model_name) + run_single(block_fn(model()), name, x) \ No newline at end of file From 8e9858787ce384dd347498d677b1b33e6ad1a31f Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Tue, 7 May 2024 21:30:56 +0530 Subject: [PATCH 04/13] more vis fixes (cherry picked from commit fb3cb8fe0dfd2f2b51570df3e36f4b219f611e00) --- experiments/util/visualize_graph.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/experiments/util/visualize_graph.py b/experiments/util/visualize_graph.py index a74a27c..79ae883 100644 --- a/experiments/util/visualize_graph.py +++ b/experiments/util/visualize_graph.py @@ -11,20 +11,20 @@ from transformers import AutoConfig to_test = { - 'bert_encoder': ['bert', lambda model: model.bert.encoder.layer[0]], - 'memsave_bert_encoder': ['memsave_bert', lambda model: model.bert.encoder.layer[0]], - 'bart_encoder': ['bart', lambda model: model.decoder.layers[0]], - 'memsave_bart_encoder': ['memsave_bart', lambda model: model.decoder.layers[0]], - 'gpt2_layer': ['gpt2', lambda model: model.transformer.h[0]], - 'memsave_gpt2_layer': ['memsave_gpt2', lambda model: model.transformer.h[0]], - 't5_decoder': ['t5', lambda model: model.decoder.block[1]], - 'memsave_t5_decoder': ['memsave_t5', lambda model: model.decoder.block[1]], + 'bert_encoder': ['bert', lambda model_full: model_full.bert.encoder.layer[0]], + 'memsave_bert_encoder': ['memsave_bert', lambda model_full: model_full.bert.encoder.layer[0]], + 'bart_encoder': ['bart', lambda model_full: model_full.model.decoder.layers[0]], + 'memsave_bart_encoder': ['memsave_bart', lambda model_full: model_full.model.decoder.layers[0]], + 'gpt2_layer': ['gpt2', lambda model_full: model_full.transformer.h[0]], + 'memsave_gpt2_layer': ['memsave_gpt2', lambda model_full: model_full.transformer.h[0]], + 't5_decoder': ['t5', lambda model_full: model_full.decoder.block[1]], + 'memsave_t5_decoder': ['memsave_t5', lambda model_full: model_full.decoder.block[1]], } def run_single(model, name, x): y = model(x) dot = make_dot( - y.mean(), + y[0].mean(), params=dict(model.named_parameters()), show_attrs=True, show_saved=True, @@ -50,7 +50,7 @@ def run_single(model, name, x): config = models.get_transformers_config(model_name) models.transformer_input_shape = (config.vocab_size, config.hidden_size) - x = torch.rand(7, *models.transformer_input_shape) + x = torch.rand(7, 128, config.hidden_size) model = models.transformer_model_fns.get(model_name) run_single(block_fn(model()), name, x) From 14f586f401894645e325ab801fec5b5ab87717c3 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:19:34 +0530 Subject: [PATCH 05/13] rms layer norm (cherry picked from commit 80438be61b847168f4a251c85604fd43252d4ab4) --- experiments/util/visualize_graph.py | 30 ++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/experiments/util/visualize_graph.py b/experiments/util/visualize_graph.py index 79ae883..2c5477a 100644 --- a/experiments/util/visualize_graph.py +++ b/experiments/util/visualize_graph.py @@ -11,16 +11,29 @@ from transformers import AutoConfig to_test = { - 'bert_encoder': ['bert', lambda model_full: model_full.bert.encoder.layer[0]], - 'memsave_bert_encoder': ['memsave_bert', lambda model_full: model_full.bert.encoder.layer[0]], - 'bart_encoder': ['bart', lambda model_full: model_full.model.decoder.layers[0]], - 'memsave_bart_encoder': ['memsave_bart', lambda model_full: model_full.model.decoder.layers[0]], - 'gpt2_layer': ['gpt2', lambda model_full: model_full.transformer.h[0]], - 'memsave_gpt2_layer': ['memsave_gpt2', lambda model_full: model_full.transformer.h[0]], - 't5_decoder': ['t5', lambda model_full: model_full.decoder.block[1]], - 'memsave_t5_decoder': ['memsave_t5', lambda model_full: model_full.decoder.block[1]], + "bert_encoder": ["bert", lambda model_full: model_full.bert.encoder.layer[0]], + "memsave_bert_encoder": [ + "memsave_bert", + lambda model_full: model_full.bert.encoder.layer[0], + ], + "bart_encoder": ["bart", lambda model_full: model_full.model.decoder.layers[0]], + "memsave_bart_encoder": [ + "memsave_bart", + lambda model_full: model_full.model.decoder.layers[0], + ], + "gpt2_layer": ["gpt2", lambda model_full: model_full.transformer.h[0]], + "memsave_gpt2_layer": [ + "memsave_gpt2", + lambda model_full: model_full.transformer.h[0], + ], + "t5_decoder": ["t5", lambda model_full: model_full.decoder.block[1]], + "memsave_t5_decoder": [ + "memsave_t5", + lambda model_full: model_full.decoder.block[1], + ], } + def run_single(model, name, x): y = model(x) dot = make_dot( @@ -54,4 +67,3 @@ def run_single(model, name, x): model = models.transformer_model_fns.get(model_name) run_single(block_fn(model()), name, x) - \ No newline at end of file From 2d997c688938b85d89db1aa6b624fd15c4175b24 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:22:13 +0530 Subject: [PATCH 06/13] code for making latex table (cherry picked from commit ec4564a93eff180b9f757f4acf342a7d0dcb6ee2) --- experiments/best_results_to_latex.py | 55 ++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 experiments/best_results_to_latex.py diff --git a/experiments/best_results_to_latex.py b/experiments/best_results_to_latex.py new file mode 100644 index 0000000..d3219e2 --- /dev/null +++ b/experiments/best_results_to_latex.py @@ -0,0 +1,55 @@ +"""Simple script to make a latex table from best results""" + +import pandas as pd + +df = pd.read_csv("results/llm/best_results-transformer-cuda-usage_stats.csv") + +df = df.set_index("model") +df = df[df["case"] != "Conv"] + +df["memsave"] = df.index.str.startswith("memsave_") +badi = df.index.map( + lambda x: x.split("memsave_", 1)[1] if x.startswith("memsave") else x +) +badi.name = "model_clean" +df2 = df.reset_index().set_index(badi).sort_index() +divs = df2[(df2["case"] == "All") & (~df2["memsave"])] +df2["Scaled M"] = df2["Memory Usage (GB)"] / divs["Memory Usage (GB)"] +df2["Scaled T"] = df2["Time Taken (s)"] / divs["Time Taken (s)"] + +df2["Memory [GiB]"] = df2.apply( + lambda x: f"{x['Memory Usage (GB)']:.2f} ({x['Scaled M']:.2f})", axis=1 +) +df2["Time [s]"] = df2.apply( + lambda x: f"{x['Time Taken (s)']:.2f} ({x['Scaled T']:.2f})", axis=1 +) + + +def _format_name(n): + if n.startswith("memsave_"): + mname = n.split("memsave_", 1)[1] + return f"{mname.capitalize()} + MemSave" + return n.capitalize() + + +ni = df2["model"].apply(_format_name) +df2 = ( + df2.set_index(ni) + .sort_index() + .drop( + columns=[ + "model", + "memsave", + "Memory Usage (GB)", + "Time Taken (s)", + "Scaled M", + "Scaled T", + ] + ) +) + +df2_p = df2.pivot_table( + index="model", columns="case", values=df2.columns[1:], aggfunc=lambda x: x +) + +print(df2_p.to_latex(na_rep="-")) From 2c0ccbae03e64972b4d66de27b1f9b60ab0ca67a Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:23:11 +0530 Subject: [PATCH 07/13] new implementation of dropout (cherry picked from commit 657e708dcac773e31a079acfe9296d86bc6234ab) --- experiments/best_results_to_latex.py | 55 +++++++++++++++++++--------- 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/experiments/best_results_to_latex.py b/experiments/best_results_to_latex.py index d3219e2..2c4c9ca 100644 --- a/experiments/best_results_to_latex.py +++ b/experiments/best_results_to_latex.py @@ -2,7 +2,13 @@ import pandas as pd -df = pd.read_csv("results/llm/best_results-transformer-cuda-usage_stats.csv") +import argparse +parser = argparse.ArgumentParser() +parser.add_argument('--input', type=str, required=True, help='Path to best_results.csv generated by get_best_results') + +args = parser.parse_args() + +df = pd.read_csv(args.input) df = df.set_index("model") df = df[df["case"] != "Conv"] @@ -24,32 +30,47 @@ lambda x: f"{x['Time Taken (s)']:.2f} ({x['Scaled T']:.2f})", axis=1 ) +def _highlight(group, col_sort, col_bold): + for c_s, c_b in zip(col_sort, col_bold): + min_idx = group[c_s].argmin() + group[c_b] = [f"\\textbf{{{group.iloc[i][c_b]}}}" if i==min_idx else group.iloc[i][c_b] for i in range(len(group.index))] + return group + +df2 = df2.groupby(['model_clean', 'case'])\ + .apply(_highlight, ['Memory Usage (GB)', 'Time Taken (s)'], ['Memory [GiB]', 'Time [s]']) + +names = { + 'bert': 'BERT', + 'bart': 'BART', + 'roberta': 'RoBERTa', + 'gpt2': 'GPT-2', + 't5': 'T5', + 'flan-t5': 'FLAN-T5', + 'mistral-7b': 'Mistral-7B', + 'transformer': 'Transformer', +} def _format_name(n): if n.startswith("memsave_"): mname = n.split("memsave_", 1)[1] - return f"{mname.capitalize()} + MemSave" - return n.capitalize() + return f"{names[mname]} + MemSave" + return names[n] ni = df2["model"].apply(_format_name) -df2 = ( - df2.set_index(ni) - .sort_index() - .drop( - columns=[ - "model", - "memsave", - "Memory Usage (GB)", - "Time Taken (s)", - "Scaled M", - "Scaled T", - ] - ) +df2 = df2.set_index(ni).sort_index().drop( + columns=[ + "model", + "memsave", + "Memory Usage (GB)", + "Time Taken (s)", + "Scaled M", + "Scaled T", + ] ) df2_p = df2.pivot_table( index="model", columns="case", values=df2.columns[1:], aggfunc=lambda x: x ) -print(df2_p.to_latex(na_rep="-")) +print(df2_p.to_latex(na_rep="-", multicolumn_format='c')) From 6e6fb005115759ceb35ce31e767f1b4f0276ebb3 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:23:59 +0530 Subject: [PATCH 08/13] adding llama3 and phi3 (cherry picked from commit 35da149c08980b8985db0238ec0a448968368d01) --- experiments/best_results_to_latex.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/experiments/best_results_to_latex.py b/experiments/best_results_to_latex.py index 2c4c9ca..4289717 100644 --- a/experiments/best_results_to_latex.py +++ b/experiments/best_results_to_latex.py @@ -37,7 +37,8 @@ def _highlight(group, col_sort, col_bold): return group df2 = df2.groupby(['model_clean', 'case'])\ - .apply(_highlight, ['Memory Usage (GB)', 'Time Taken (s)'], ['Memory [GiB]', 'Time [s]']) + .apply(_highlight, ['Memory Usage (GB)'], ['Memory [GiB]']) + # .apply(_highlight, ['Memory Usage (GB)', 'Time Taken (s)'], ['Memory [GiB]', 'Time [s]']) names = { 'bert': 'BERT', From 4568c03f167b1f5c01211ae3897319a43f85bc16 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:24:39 +0530 Subject: [PATCH 09/13] add VLM architecture in estimate, refactor cases to refer to a central location (cherry picked from commit 1f9fb45f84c12e20ef13e33bb1766a509990289b) --- experiments/best_results_to_latex.py | 48 +++++++++++++++++++--------- 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/experiments/best_results_to_latex.py b/experiments/best_results_to_latex.py index 4289717..a2e4c5f 100644 --- a/experiments/best_results_to_latex.py +++ b/experiments/best_results_to_latex.py @@ -3,8 +3,14 @@ import pandas as pd import argparse + parser = argparse.ArgumentParser() -parser.add_argument('--input', type=str, required=True, help='Path to best_results.csv generated by get_best_results') +parser.add_argument( + "--input", + type=str, + required=True, + help="Path to best_results.csv generated by get_best_results", +) args = parser.parse_args() @@ -30,27 +36,36 @@ lambda x: f"{x['Time Taken (s)']:.2f} ({x['Scaled T']:.2f})", axis=1 ) + def _highlight(group, col_sort, col_bold): for c_s, c_b in zip(col_sort, col_bold): min_idx = group[c_s].argmin() - group[c_b] = [f"\\textbf{{{group.iloc[i][c_b]}}}" if i==min_idx else group.iloc[i][c_b] for i in range(len(group.index))] + group[c_b] = [ + f"\\textbf{{{group.iloc[i][c_b]}}}" if i == min_idx else group.iloc[i][c_b] + for i in range(len(group.index)) + ] return group -df2 = df2.groupby(['model_clean', 'case'])\ - .apply(_highlight, ['Memory Usage (GB)'], ['Memory [GiB]']) - # .apply(_highlight, ['Memory Usage (GB)', 'Time Taken (s)'], ['Memory [GiB]', 'Time [s]']) + +df2 = df2.groupby(["model_clean", "case"]).apply( + _highlight, ["Memory Usage (GB)"], ["Memory [GiB]"] +) +# .apply(_highlight, ['Memory Usage (GB)', 'Time Taken (s)'], ['Memory [GiB]', 'Time [s]']) names = { - 'bert': 'BERT', - 'bart': 'BART', - 'roberta': 'RoBERTa', - 'gpt2': 'GPT-2', - 't5': 'T5', - 'flan-t5': 'FLAN-T5', - 'mistral-7b': 'Mistral-7B', - 'transformer': 'Transformer', + "bert": "BERT", + "bart": "BART", + "roberta": "RoBERTa", + "gpt2": "GPT-2", + "t5": "T5", + "flan-t5": "FLAN-T5", + "mistral-7b": "Mistral-7B", + "transformer": "Transformer", + "llama3-8b": "LLaMa3-8B", + "phi3-4b": "Phi3-4B", } + def _format_name(n): if n.startswith("memsave_"): mname = n.split("memsave_", 1)[1] @@ -68,10 +83,13 @@ def _format_name(n): "Scaled M", "Scaled T", ] -) +) # fmt: skip df2_p = df2.pivot_table( index="model", columns="case", values=df2.columns[1:], aggfunc=lambda x: x ) -print(df2_p.to_latex(na_rep="-", multicolumn_format='c')) +short_index = df2_p.index.map(lambda t: "+ MemSave" if "+ MemSave" in t else t) +df2_p = df2_p.set_index(short_index) + +print(df2_p.to_latex(na_rep="-", multicolumn_format="c")) From 6c26608f463b31163f55aa80ac4ea92ec64b866e Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Sun, 2 Jun 2024 02:29:02 +0530 Subject: [PATCH 10/13] torchviz visualize elements (cherry picked from commit e9b686ff937a1f16b9ee098a950c775382489f1e) --- .../util/visualize_graph_elementary.py | 142 ++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 experiments/util/visualize_graph_elementary.py diff --git a/experiments/util/visualize_graph_elementary.py b/experiments/util/visualize_graph_elementary.py new file mode 100644 index 0000000..2d4001a --- /dev/null +++ b/experiments/util/visualize_graph_elementary.py @@ -0,0 +1,142 @@ +# ruff: noqa +import argparse + +import torch +from torchview import draw_graph +from torchviz import make_dot + +from tqdm import tqdm +import memsave_torch +from experiments.util.collect_results import select_cases +from experiments.util.estimate import parse_case +from experiments.util.measurements import separate_grad_arguments + +to_test = [ + { + "name": "Linear2dims", + "layer_fn": lambda: torch.nn.Linear(3, 5), + "data_fn": lambda: torch.rand(7, 12, 3), # weight sharing + "grads": ["All", "Input", "Linear"], + }, + { + "name": "Conv2d", + "layer_fn": lambda: torch.nn.Conv2d(3, 5, 3), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input", "Conv"], + }, + { + "name": "BatchNorm2d", + "layer_fn": lambda: torch.nn.BatchNorm2d(3), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input", "Norm"], + }, + { + "name": "LayerNorm", + "layer_fn": lambda: torch.nn.LayerNorm([3, 12, 12]), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input", "Norm"], + }, + { + "name": "Dropout", + "layer_fn": lambda: torch.nn.Dropout(), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input"], + }, + { + "name": "MaxPool2d", + "layer_fn": lambda: torch.nn.MaxPool2d(3), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input"], + }, + { + "name": "ReLU", + "layer_fn": lambda: torch.nn.ReLU(), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input"], + }, + { + "name": "SiLU", + "layer_fn": lambda: torch.nn.SiLU(), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input"], + }, +] + + +def run_single(model, x, name, dirname): + y = model(x) + dot = make_dot( + y.sum(), + params=dict(model.named_parameters()), + show_attrs=True, + show_saved=True, + ) + dot.render(filename=name, directory=dirname) + + +def separate_grad_arguments_wrapper( + model, + grad_linear_weights: bool = True, + grad_linear_bias: bool = True, + grad_conv_weights: bool = True, + grad_conv_bias: bool = True, + grad_norm_weights: bool = True, + grad_norm_bias: bool = True, + grad_embed_weights: bool = False, + **kwargs, +): + return separate_grad_arguments( + model, + grad_linear_weights=grad_linear_weights, + grad_linear_bias=grad_linear_bias, + grad_conv_weights=grad_conv_weights, + grad_conv_bias=grad_conv_bias, + grad_norm_weights=grad_norm_weights, + grad_norm_bias=grad_norm_bias, + grad_embed_weights=grad_embed_weights, + ) + + +if __name__ == "__main__": + for layer_to_test in (pbar := tqdm(to_test)): + pbar.set_description(layer_to_test["name"]) + all_grad_cases = select_cases(layer_to_test["grads"]) + for c_name, c in zip(layer_to_test["grads"], all_grad_cases): + grad_opts = parse_case(c) + x = layer_to_test["data_fn"]() + layer = layer_to_test["layer_fn"]() + memsave_layer = memsave_torch.nn.convert_to_memory_saving( + layer, clone_params=True + ) + leafs, no_leafs = separate_grad_arguments_wrapper( + layer, **grad_opts + ) # no weights differentiable + + grad_input = False + x2 = x.clone() + grad_input = "grad_input" in grad_opts and grad_opts["grad_input"] + + leafs = ([x] if grad_input else []) + leafs + no_leafs = ([] if grad_input else [x]) + no_leafs + + for leaf in leafs: + leaf.requires_grad_(True) + for no_leaf in no_leafs: + no_leaf.requires_grad_(False) + + # TODO: add grad weights case + leafs, no_leafs = separate_grad_arguments_wrapper( + memsave_layer, **grad_opts + ) # no weights differentiable + + leafs = ([x2] if grad_input else []) + leafs + no_leafs = ([] if grad_input else [x2]) + no_leafs + + for leaf in leafs: + leaf.requires_grad_(True) + for no_leaf in no_leafs: + no_leaf.requires_grad_(False) + + dirname = f"torchviz-output/elementary/{layer_to_test['name']}" + run_single(layer, x, c_name, dirname) + run_single(memsave_layer, x2, c_name + "_MemSave", dirname) From bdf34afd54796da1d6055bae17ce317143d19b1b Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:25:26 +0530 Subject: [PATCH 11/13] format + prep for merge (cherry picked from commit c444edd2c1532dd6b1cfb6120018a46e51666a25) --- .gitignore | 2 ++ experiments/best_results_to_latex.py | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 6ad7868..7f71819 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ *.txt *.csv !requirements.txt +torchviz-output/ +torchview-output/ # generated docs docs_src/_build/ diff --git a/experiments/best_results_to_latex.py b/experiments/best_results_to_latex.py index a2e4c5f..c62ef58 100644 --- a/experiments/best_results_to_latex.py +++ b/experiments/best_results_to_latex.py @@ -1,9 +1,9 @@ """Simple script to make a latex table from best results""" -import pandas as pd - import argparse +import pandas as pd + parser = argparse.ArgumentParser() parser.add_argument( "--input", From 8eed5476b51d8aaf307acc52d0acc722c33799f4 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 13:27:28 +0530 Subject: [PATCH 12/13] minor improvements to latex table script and batch norm in eval (cherry picked from commit 8e68a88d4e079450eac350691d0d8e14648fbf9c) --- experiments/best_results_to_latex.py | 33 +++++++++++++++---- .../util/visualize_graph_elementary.py | 29 +++++++++++----- 2 files changed, 48 insertions(+), 14 deletions(-) diff --git a/experiments/best_results_to_latex.py b/experiments/best_results_to_latex.py index c62ef58..4b4e745 100644 --- a/experiments/best_results_to_latex.py +++ b/experiments/best_results_to_latex.py @@ -12,12 +12,14 @@ help="Path to best_results.csv generated by get_best_results", ) + args = parser.parse_args() df = pd.read_csv(args.input) df = df.set_index("model") df = df[df["case"] != "Conv"] +df = df[df["case"] != "SurgicalLast"] df["memsave"] = df.index.str.startswith("memsave_") badi = df.index.map( @@ -57,14 +59,26 @@ def _highlight(group, col_sort, col_bold): "bart": "BART", "roberta": "RoBERTa", "gpt2": "GPT-2", - "t5": "T5", + "t5": "T5 \\cite{JMLR_t5}", "flan-t5": "FLAN-T5", - "mistral-7b": "Mistral-7B", - "transformer": "Transformer", - "llama3-8b": "LLaMa3-8B", - "phi3-4b": "Phi3-4B", + "mistral-7b": "Mistral-7B \\cite{jiang2023mistral}", + "transformer": "Transformer \\cite{NIPS2017_3f5ee243_vaswaniattention}", + "llama3-8b": "LLaMa3-8B \\cite{touvron2023llama}", + "phi3-4b": "Phi3-4B \\cite{gunasekar2023textbooksPhi}", + # Conv + "deeplabv3_resnet101": "DeepLabv3 (RN101) \\cite{deeplabv3_chen2017rethinking}", + "efficientnet_v2_l": "EfficientNetv2-L \\cite{efficientnet_TanL19,efficientnetv2_TanL21}", + "fcn_resnet101": "FCN (RN101) \\cite{fcn}", + "mobilenet_v3_large": "MobileNetv3-L \\cite{mobilenetv3}", + "resnext101_64x4d": "ResNeXt101-64x4d \\cite{resnext_cvpr_XieGDTH17}", + "fasterrcnn_resnet50_fpn_v2": "Faster-RCNN (RN101) \\cite{faster_rcnn_RenHGS15}", + "ssdlite320_mobilenet_v3_large": "SSDLite (MobileNetv3-L) \\cite{mobilenetv2_Sandler_2018_CVPR}", + "vgg16": "VGG-16 \\cite{vgg_SimonyanZ14a}", } +# import ipdb; ipdb.set_trace() +df2 = df2[df2.index.isin(names.keys(), level=0)] + def _format_name(n): if n.startswith("memsave_"): @@ -92,4 +106,11 @@ def _format_name(n): short_index = df2_p.index.map(lambda t: "+ MemSave" if "+ MemSave" in t else t) df2_p = df2_p.set_index(short_index) -print(df2_p.to_latex(na_rep="-", multicolumn_format="c")) +latex_str = df2_p.to_latex(na_rep="-", multicolumn_format="c") +final_str = "" +for line in latex_str.split("\n"): + add_line = line + "\n" + if line.startswith("+ MemSave"): + add_line += "\\midrule\n" + final_str += add_line +print(final_str) diff --git a/experiments/util/visualize_graph_elementary.py b/experiments/util/visualize_graph_elementary.py index 2d4001a..642bd24 100644 --- a/experiments/util/visualize_graph_elementary.py +++ b/experiments/util/visualize_graph_elementary.py @@ -11,54 +11,67 @@ from experiments.util.estimate import parse_case from experiments.util.measurements import separate_grad_arguments + +def eval_bn(num_features): + m = torch.nn.BatchNorm2d(num_features) + m.eval() + return m + + to_test = [ { "name": "Linear2dims", "layer_fn": lambda: torch.nn.Linear(3, 5), "data_fn": lambda: torch.rand(7, 12, 3), # weight sharing - "grads": ["All", "Input", "Linear"], + "grads": ["All", "Input", "Linear", "Everything", "Nothing"], }, { "name": "Conv2d", "layer_fn": lambda: torch.nn.Conv2d(3, 5, 3), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input", "Conv"], + "grads": ["All", "Input", "Conv", "Everything", "Nothing"], }, { "name": "BatchNorm2d", "layer_fn": lambda: torch.nn.BatchNorm2d(3), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input", "Norm"], + "grads": ["All", "Input", "Norm", "Everything", "Nothing"], + }, + { + "name": "BatchNorm2d_Eval", + "layer_fn": lambda: eval_bn(3), + "data_fn": lambda: torch.rand(7, 3, 12, 12), + "grads": ["All", "Input", "Norm", "Everything", "Nothing"], }, { "name": "LayerNorm", "layer_fn": lambda: torch.nn.LayerNorm([3, 12, 12]), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input", "Norm"], + "grads": ["All", "Input", "Norm", "Everything", "Nothing"], }, { "name": "Dropout", "layer_fn": lambda: torch.nn.Dropout(), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input"], + "grads": ["All", "Input", "Everything", "Nothing"], }, { "name": "MaxPool2d", "layer_fn": lambda: torch.nn.MaxPool2d(3), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input"], + "grads": ["All", "Input", "Everything", "Nothing"], }, { "name": "ReLU", "layer_fn": lambda: torch.nn.ReLU(), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input"], + "grads": ["All", "Input", "Everything", "Nothing"], }, { "name": "SiLU", "layer_fn": lambda: torch.nn.SiLU(), "data_fn": lambda: torch.rand(7, 3, 12, 12), - "grads": ["All", "Input"], + "grads": ["All", "Input", "Everything", "Nothing"], }, ] From 22957d90cb6e907718d789d6d124c5b342ebcf93 Mon Sep 17 00:00:00 2001 From: Plutonium-239 Date: Thu, 22 Aug 2024 14:17:12 +0530 Subject: [PATCH 13/13] minor updates: - add script to convert resnet best results to latex table - formatting to make ruff happy (cherry picked from commit f69f725f1423788303e938ab47b2c0882aece5e0) --- experiments/resnet_best_results_to_latex.py | 39 +++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 experiments/resnet_best_results_to_latex.py diff --git a/experiments/resnet_best_results_to_latex.py b/experiments/resnet_best_results_to_latex.py new file mode 100644 index 0000000..c9d4f77 --- /dev/null +++ b/experiments/resnet_best_results_to_latex.py @@ -0,0 +1,39 @@ +"""Simple script to make a latex table from resnet results""" + +import pandas as pd + +df = pd.read_csv("results/resnet101_only/best_results-conv-cpu-usage_stats.csv") +df = df.set_index("model") +df = df.drop(columns=["Scaled M", "Scaled T"]) +df = df.drop("memsave_resnet101_conv+relu+bn") +df = df[df["case"] != "SurgicalLast"] +df = df[df["case"] != "Conv"] + +mem_div = df[df["case"] == "All"].loc["resnet101", "Memory Usage (GB)"] +time_div = df[df["case"] == "All"].loc["resnet101", "Time Taken (s)"] +df["Scaled M"] = df["Memory Usage (GB)"] / mem_div +df["Scaled T"] = df["Time Taken (s)"] / time_div + +df["Memory [GiB]"] = df.apply( + lambda x: f"{x['Memory Usage (GB)']:.2f} ({x['Scaled M']:.2f})", axis=1 +) +df["Time [s]"] = df.apply( + lambda x: f"{x['Time Taken (s)']:.2f} ({x['Scaled T']:.2f})", axis=1 +) + +df = df.drop(columns=["Scaled M", "Scaled T", "Memory Usage (GB)", "Time Taken (s)"]) +df_p = df.pivot_table( + index="model", columns="case", values=df.columns[1:], aggfunc=lambda x: x +) + +labels = { + "resnet101": "Default ResNet-101", + "memsave_resnet101_conv": "+ swap Convolution", + "memsave_resnet101_conv_full": "+ swap BatchNorm, ReLU", +} + +df_p = df_p.rename(index=labels) +df_p = df_p.sort_index(ascending=False) + +print(df_p["Memory [GiB]"].to_latex()) +print(df_p["Time [s]"].to_latex())