| |
|
|
| library(tidyverse) |
| library(dplyr) |
| library(readr) |
| library(ggthemes) |
| library(ggplot2) |
| library(patchwork) |
| library(optparse) |
| library(RColorBrewer) |
|
|
| plot_allot_pair <- function(df_lower, df_upper, disp_fnorm = TRUE) { |
| ylabel <- ifelse(disp_fnorm == TRUE, "FNorm", "Sensitivity") |
| mem_plot <- ggplot( |
| df_upper, aes(x = layer, y = memmb, color = "blue") |
| ) + |
| geom_line() + |
| theme_gray(base_size = 12) + |
| labs(x = "", y = "Memory(MiB)") + |
| theme( |
| legend.position = "none", |
| axis.title.x = element_blank(), |
| axis.text.x = element_blank(), |
| axis.ticks = element_blank() |
| ) + |
| scale_color_solarized() |
|
|
| metric_cutoff <- 1.0 |
| if (is.na(df_lower$attempt[1])) { |
| module_disp <- paste0("HQQ - ", df_lower$module[1]) |
| } else { |
| module_disp <- paste0(df_lower$attempt[1], " - ", df_lower$module[1]) |
| } |
| bar_plot <- ggplot( |
| df_lower, aes(x = layer, y = metric, fill = cfg, color = cfg) |
| ) + |
| geom_bar(stat = "identity", color = "gray50") + |
| geom_text( |
| data = subset(df_lower, metric <= metric_cutoff), |
| aes(x = layer, label = cfg), |
| colour = "black", |
| angle = 90, |
| hjust = -0.20, |
| size = 2 |
| ) + |
| geom_text( |
| data = subset(df_lower, metric > metric_cutoff), |
| aes(x = layer, label = cfg), |
| colour = "black", |
| angle = 90, |
| vjust = 0.20, |
| position = position_stack(vjust = 0.5), |
| size = 2 |
| ) + |
| theme_gray(base_size = 12) + |
| labs(x = module_disp, y = "FNorm") + |
| theme(legend.position = "none") + |
| col_scale |
|
|
| |
| combined_plot <- mem_plot / bar_plot + plot_layout(heights = c(1, 2)) |
| return(combined_plot) |
| } |
|
|
| weight_grid <- function(df_disp, mod, attmpt, bar_value_fnorm) { |
| if (is.null(attmpt)) { |
| df_plot <- df_disp |> filter(module == mod & is.na(attempt)) |
| } else { |
| df_plot <- df_disp |> filter(module == mod & attempt == attmpt) |
| } |
| plot_allot_pair(df_plot, df_plot, bar_value_fnorm) |
| } |
|
|
| strip_name <- function(name) { |
| start <- nchar("fnorm-") + 1 |
| stop <- nchar(name) - 4 |
| return(substr(name, start, stop)) |
| } |
|
|
| match_hqq_base_config <- function(budget) { |
| buckets <- c( |
| 2.13, 2.25, 2.51, 3.13, 3.25, 3.51, 4.13, 4.25, 4.51, 8.13, 8.25, 8.51 |
| ) |
| idx <- which.min(abs(buckets - budget)) |
| return(buckets[[idx]]) |
| } |
|
|
| parser <- OptionParser() |
| parser <- add_option( |
| parser, c("-m", "--model"), |
| type = "character", |
| help = "Model ID", |
| metavar = "character" |
| ) |
| parser <- add_option( |
| parser, c("-b", "--budget"), |
| type = "double", |
| help = "Bit Budget", |
| metavar = "double" |
| ) |
| parser <- add_option( |
| parser, c("-d", "--fnorm_data_dir"), |
| type = "character", |
| help = "Data directory of fnorm meta data", |
| metavar = "character" |
| ) |
| parser <- add_option( |
| parser, c("-q", "--quant_cfg_allot_file"), |
| type = "character", |
| help = "The combined quant config allocation csv file", |
| metavar = "character" |
| ) |
| parser <- add_option( |
| parser, c("--attempt1"), |
| type = "character", |
| help = "The first attempt to plot", |
| metavar = "character" |
| ) |
| parser <- add_option( |
| parser, c("--attempt2"), |
| type = "character", |
| help = "The second attempt to plot", |
| metavar = "character" |
| ) |
| parser <- add_option( |
| parser, c("--fnorm"), |
| action = "store_true", |
| default = TRUE, |
| type = "logical", |
| help = "Display FNorm value in the bar chart", |
| metavar = "logical" |
| ) |
|
|
|
|
| args <- parse_args(parser) |
|
|
| if (is.null(args$model)) { |
| model_id <- "Llama-2-7b-hf" |
| } else { |
| model_id <- args$model |
| } |
|
|
| if (is.null(args$budget)) { |
| budget <- 4.51 |
| } else { |
| budget <- args$budget |
| } |
|
|
| if (is.null(args$fnorm_data_dir)) { |
| fnorm_data_dir <- "../src/data" |
| } else { |
| fnorm_data_dir <- args$fnorm_data_dir |
| } |
|
|
| if (is.null(args$quant_cfg_allot_file)) { |
| quant_cfg_allot_file <- "data/quant-cfg-allocation.csv" |
| } else { |
| quant_cfg_allot_file <- args$quant_cfg_allot_file |
| } |
|
|
| if (is.null(args$output_dir)) { |
| output_dir <- "pdfs/allot" |
| } else { |
| output_dir <- args$output_dir |
| } |
|
|
| if (is.null(args$attempt1)) { |
| attempt1 <- "mxq1" |
| } else { |
| attempt1 <- args$attempt1 |
| } |
|
|
| if (is.null(args$attempt2)) { |
| attempt2 <- "kurt-scaled" |
| } else { |
| attempt2 <- args$attempt2 |
| } |
|
|
| if (is.null(args$fnorm)) { |
| bar_value_fnorm <- FALSE |
| } else { |
| bar_value_fnorm <- args$fnorm |
| } |
|
|
| fnorm_dir <- path.expand(fnorm_data_dir) |
| fnorm_fps <- dir( |
| path = fnorm_dir, |
| pattern = paste0("fnorm-", model_id, "\\.csv$"), |
| |
| full.names = TRUE |
| ) |
| names(fnorm_fps) <- sapply((basename(fnorm_fps)), strip_name) |
| df_fnorm <- plyr::ldply( |
| fnorm_fps, |
| read.csv, |
| stringsAsFactors = FALSE, |
| .id = "model" |
| ) |
|
|
| k_cols <- c( |
| "model", "module", "layer", "cfg", "fnorm", "sensitivity", "memmb", "params" |
| ) |
| df_fnorm <- df_fnorm |> |
| filter( |
| model == model_id |
| ) |> |
| dplyr::mutate( |
| cfg = paste0("b", nbit1, "g", gsize1) |
| ) |> |
| select(all_of(k_cols)) |> |
| dplyr::mutate( |
| cfg = factor( |
| cfg, |
| levels = c( |
| "b2g128", "b2g64", "b2g32", |
| "b3g128", "b3g64", "b3g32", |
| "b4g128", "b4g64", "b4g32", |
| "b8g128", "b8g64", "b8g32" |
| ) |
| ) |
| ) |
|
|
| if (bar_value_fnorm) { |
| df_fnorm <- df_fnorm |> |
| dplyr::rename(metric = fnorm) |> |
| select(-c("sensitivity")) |
| } else { |
| df_fnorm <- df_fnorm |> |
| dplyr::rename(metric = sensitivity) |> |
| select(-c("fnorm")) |
| } |
|
|
| rounded_budget <- match_hqq_base_config(budget) |
| df_cfgs <- read_csv("data/quant-cfg-allocation.csv") |
|
|
| df_cfg_1 <- df_cfgs |> |
| filter( |
| |
| bit_budget == budget & model == model_id |
| ) |> |
| dplyr::mutate( |
| cfg = paste0("b", b1, "g", g1) |
| ) |> |
| select(-c("b1", "g1", "b2", "g2", "bit_budget", "memmb")) |> |
| dplyr::mutate( |
| cfg = factor( |
| cfg, |
| levels = c( |
| "b2g128", "b2g64", "b2g32", |
| "b3g128", "b3g64", "b3g32", |
| "b4g128", "b4g64", "b4g32", |
| "b8g128", "b8g64", "b8g32" |
| ) |
| ) |
| ) |
|
|
| by <- join_by(model == model, module == module, layer == layer, cfg == cfg) |
| df_disp <- df_cfg_1 |> |
| left_join(df_fnorm, by) |
|
|
| my_colors <- brewer.pal(12, "Paired") |
| names(my_colors) <- levels(df_disp$cfg) |
| col_scale <- scale_fill_manual(name = "cfg", values = my_colors) |
|
|
| p11 <- weight_grid(df_disp, "mlp.down_proj", NULL, bar_value_fnorm) |
| p12 <- weight_grid(df_disp, "mlp.down_proj", attempt1, bar_value_fnorm) |
| p13 <- weight_grid(df_disp, "mlp.down_proj", attempt2, bar_value_fnorm) |
|
|
| p21 <- weight_grid(df_disp, "mlp.gate_proj", NULL, bar_value_fnorm) |
| p22 <- weight_grid(df_disp, "mlp.gate_proj", attempt1, bar_value_fnorm) |
| p23 <- weight_grid(df_disp, "mlp.gate_proj", attempt2, bar_value_fnorm) |
|
|
| p31 <- weight_grid(df_disp, "mlp.up_proj", NULL, bar_value_fnorm) |
| p32 <- weight_grid(df_disp, "mlp.up_proj", attempt1, bar_value_fnorm) |
| p33 <- weight_grid(df_disp, "mlp.up_proj", attempt2, bar_value_fnorm) |
|
|
| final_plot <- (p11 | p21 | p31) / (p12 | p22 | p32) / (p13 | p23 | p33) |
| final_plot |
| ggsave( |
| paste0(output_dir, "/", model_id, "-", budget, "-mlp-allot.pdf"), |
| create.dir = TRUE, |
| width = 16, height = 9 |
| ) |
|
|
| p61 <- weight_grid(df_disp, "self_attn.k_proj", NULL, bar_value_fnorm) |
| p62 <- weight_grid(df_disp, "self_attn.k_proj", attempt1, bar_value_fnorm) |
| p63 <- weight_grid(df_disp, "self_attn.k_proj", attempt2, bar_value_fnorm) |
|
|
| p71 <- weight_grid(df_disp, "self_attn.o_proj", NULL, bar_value_fnorm) |
| p72 <- weight_grid(df_disp, "self_attn.o_proj", attempt1, bar_value_fnorm) |
| p73 <- weight_grid(df_disp, "self_attn.o_proj", attempt2, bar_value_fnorm) |
|
|
| p81 <- weight_grid(df_disp, "self_attn.q_proj", NULL, bar_value_fnorm) |
| p82 <- weight_grid(df_disp, "self_attn.q_proj", attempt1, bar_value_fnorm) |
| p83 <- weight_grid(df_disp, "self_attn.q_proj", attempt2, bar_value_fnorm) |
|
|
| p91 <- weight_grid(df_disp, "self_attn.v_proj", NULL, bar_value_fnorm) |
| p92 <- weight_grid(df_disp, "self_attn.v_proj", attempt1, bar_value_fnorm) |
| p93 <- weight_grid(df_disp, "self_attn.v_proj", attempt2, bar_value_fnorm) |
|
|
| final_plot2 <- (p61 | p71 | p81 | p91) / (p62 | p72 | p82 | p92) / (p63 | p73 | p83 | p93) |
| final_plot2 |
| ggsave( |
| paste0(output_dir, "/", model_id, "-", budget, "-attn-allot.pdf"), |
| create.dir = TRUE, |
| width = 16, height = 9 |
| ) |
|
|
