File size: 6,575 Bytes
21ad80b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | library(tidyverse)
library(readr)
library(ggthemes)
library(ggplot2)
library(patchwork)
weight_grid <- function(
df_wdist, df_kurtosis, mod, show_legend = FALSE, show_cfg = TRUE) {
df_mod_wdist <- df_wdist |> filter(module == mod)
df_mod_kurt <- df_kurtosis |> filter(module == mod)
# Line plot (on top)
line_plot <- ggplot(df_mod_kurt, aes(x = layer, y = kurtosis)) +
geom_line(color = "blue") +
theme_gray(base_size = 14) +
theme_minimal() +
theme(
axis.title.x = element_blank(),
axis.text.x = element_blank()
)
# Bar plot (on bottom)
module_disp <- df_mod_wdist$mod_disp[1]
bar_plot <- ggplot(
df_mod_wdist, aes(x = layer, y = abs_val, fill = nth_percentile)
) +
geom_bar(stat = "identity", color = "gray50") +
theme_gray(base_size = 14) +
labs(
x = module_disp, y = "Absolute Value", fill = "nth percentile"
)
if (show_cfg) {
bar_plot <- bar_plot +
geom_text(
data = subset(df_mod_wdist, nth_percentile == 100),
aes(x = layer, label = quant_cfg),
angle = 90,
vjust = 0.20,
position = position_stack(vjust = 0.5),
colour = "white",
size = 2
)
}
if (show_legend) {
bar_plot <- bar_plot +
theme(
legend.position = "bottom",
legend.text = element_text(size = 16),
legend.title = element_text(size = 16)
) +
# coord_flip() +
scale_color_solarized()
} else {
bar_plot <- bar_plot +
theme(legend.position = "none") +
scale_color_solarized()
}
# Combine the line and bar plot vertically
combined_plot <- line_plot / bar_plot + plot_layout(heights = c(1, 3))
return(combined_plot)
}
weight_grid_only <- function(
df_wdist, df_kurtosis, mod, show_legend = FALSE) {
return(weight_grid(df_wdist, df_kurtosis, mod, show_legend, show_cfg = FALSE))
}
plot_quant_cfg <- function(model_id, budget, attempt, cfg_csv_fp) {
df_cfg_all <- read_csv(cfg_csv_fp)
df_cfg <- df_cfg_all |>
filter(bit_budget == budget) |>
mutate(
quant_cfg = paste0("b", b1, "g", g1)
) |>
select(-c("b1", "g1", "b2", "g2", "bit_budget"))
percentiles <- c("0", "99", "99.9", "99.99", "100")
all_cols <- c("module", "layer", percentiles)
df_wdist <- df_all |>
mutate(
`0` = percentile_0,
`99` = percentile_99 - percentile_0,
`99.9` = percentile_999 - percentile_99,
`99.99` = percentile_9999 - percentile_999,
`100` = percentile_100 - percentile_9999,
) |>
select(all_of(all_cols)) |>
pivot_longer(
cols = percentiles,
names_to = "nth_percentile",
names_transform = list(nth_percentile = as.numeric),
values_to = "abs_val"
) |>
mutate(
nth_percentile = factor(nth_percentile, levels = rev(percentiles))
) |>
left_join(df_module_param_count, by = c("module")) |>
left_join(df_cfg, by = c("module", "layer"))
k_cols <- c("module", "layer", "kurtosis")
df_kurtosis <- df_all |>
select(all_of(k_cols))
p1 <- weight_grid(df_wdist, df_kurtosis, "input_layernorm")
p2 <- weight_grid(df_wdist, df_kurtosis, "mlp.down_proj")
p3 <- weight_grid(df_wdist, df_kurtosis, "mlp.gate_proj")
p4 <- weight_grid(df_wdist, df_kurtosis, "mlp.up_proj")
p5 <- weight_grid(df_wdist, df_kurtosis, "post_attention_layernorm")
p6 <- weight_grid(df_wdist, df_kurtosis, "self_attn.k_proj")
p7 <- weight_grid(df_wdist, df_kurtosis, "self_attn.o_proj")
p8 <- weight_grid(df_wdist, df_kurtosis, "self_attn.q_proj", TRUE)
p9 <- weight_grid(df_wdist, df_kurtosis, "self_attn.v_proj")
# Create a 3x3 grid of combined plots
final_plot <- (p1 | p2 | p3) / (p4 | p5 | p6) / (p7 | p8 | p9)
ggsave(
paste0("pdfs/", model_id, "-mxq-cfgs-from-model-", attempt, ".pdf"),
width = 16,
height = 9
)
return(final_plot)
}
plot_wdist <- function(model_id, budget, cfg_csv_fp) {
df_cfg_all <- read_csv(cfg_csv_fp)
df_cfg <- df_cfg_all |>
filter(bit_budget == budget) |>
mutate(
quant_cfg = paste0("b", b1, "g", g1)
) |>
select(-c("b1", "g1", "b2", "g2", "bit_budget"))
percentiles <- c("0", "99", "99.9", "99.99", "100")
all_cols <- c("module", "layer", percentiles)
df_wdist <- df_all |>
mutate(
`0` = percentile_0,
`99` = percentile_99 - percentile_0,
`99.9` = percentile_999 - percentile_99,
`99.99` = percentile_9999 - percentile_999,
`100` = percentile_100 - percentile_9999,
) |>
select(all_of(all_cols)) |>
pivot_longer(
cols = percentiles,
names_to = "nth_percentile",
names_transform = list(nth_percentile = as.numeric),
values_to = "abs_val"
) |>
mutate(
nth_percentile = factor(nth_percentile, levels = rev(percentiles))
) |>
left_join(df_module_param_count, by = c("module")) |>
left_join(df_cfg, by = c("module", "layer"))
k_cols <- c("module", "layer", "kurtosis")
df_kurtosis <- df_all |>
select(all_of(k_cols))
p1 <- weight_grid_only(df_wdist, df_kurtosis, "input_layernorm")
p2 <- weight_grid_only(df_wdist, df_kurtosis, "mlp.down_proj")
p3 <- weight_grid_only(df_wdist, df_kurtosis, "mlp.gate_proj")
p4 <- weight_grid_only(df_wdist, df_kurtosis, "mlp.up_proj")
p5 <- weight_grid_only(df_wdist, df_kurtosis, "post_attention_layernorm")
p6 <- weight_grid_only(df_wdist, df_kurtosis, "self_attn.k_proj")
p7 <- weight_grid_only(df_wdist, df_kurtosis, "self_attn.o_proj")
p8 <- weight_grid_only(df_wdist, df_kurtosis, "self_attn.q_proj", TRUE)
p9 <- weight_grid_only(df_wdist, df_kurtosis, "self_attn.v_proj")
# Create a 3x3 grid of combined plots
final_plot <- (p1 | p2 | p3) / (p4 | p5 | p6) / (p7 | p8 | p9)
ggsave(
paste0("pdfs/", model_id, "-wdist-kurtosis.pdf"),
width = 16, height = 9
)
return(final_plot)
}
model_id <- "Llama-2-13b-hf"
df_all <- read_csv(paste0("data/wdist/wdist-", model_id, ".csv"))
df_module_param_count <- df_all |>
select(
module, param_count
) |>
group_by(module) |>
summarise(
param_count = sum(param_count)
) |>
mutate(
mod_disp = paste0(module, "(", formatC(param_count, big.mark = ","), ")")
)
budget <- 4.51
attempt <- "MXQ1"
cfg_csv_fp <- "data/llama-mxq-cfgs.csv"
plot_quant_cfg(model_id, budget, attempt, cfg_csv_fp)
attempt <- "kurt-global"
cfg_csv_fp <- "data/kurt/global/llama-mxq-cfgs.csv"
plot_quant_cfg(model_id, budget, attempt, cfg_csv_fp)
attempt <- "kurt-scaled"
cfg_csv_fp <- "data/kurt/scaled/llama-mxq-cfgs.csv"
plot_quant_cfg(model_id, budget, attempt, cfg_csv_fp)
|