import os

import torch

import json

import string

import time

import numpy as np

import builtins as __builtin__

import io

try:

from .permutation_search_kernels import (

accelerated_search_for_good_permutation,

sum_after_2_to_4,

)

print("[ASP][Info] permutation_search_kernels can be imported.")

except ImportError:

print("[ASP][Warning] permutation_search_kernels cannot be imported.")

print(

"[ASP][Warning] If you want to accelerate the permutation search process by GPU, please build APEX by following the instructions at https://github.com/NVIDIA/apex/blob/master/apex/contrib/sparsity/README.md"

)

def convert_fx_node_name(fx_node_name):

"""Standardize punctuation of a node's name: replace all '_' with '.'"""

return fx_node_name.replace("_", ".")

def get_node_parent_children(fx_node):

"""Populate lists of all direct parents and children of a node"""

# get node parent list, and convert node name to module name

node_parent_name_converted = []

if len(fx_node.all_input_nodes) > 0:

node_parent = fx_node.all_input_nodes

for item in node_parent:

converted_item = convert_fx_node_name(item.name)

node_parent_name_converted.append(converted_item)

else:

node_parent = []

# get node children list, and convert node name to module name

node_children_name_converted = []

if len(list(fx_node.users.keys())) > 0:

node_children = list(fx_node.users.keys())

for item in node_children:

converted_item = convert_fx_node_name(item.name)

node_children_name_converted.append(converted_item)

else:

node_children = []

return node_parent_name_converted, node_children_name_converted

def node_name_matches(node_name, module_name):

"""Check for a match between graph node name and stored module name, accounting for formatting and DDP training differences"""

# process: remove all punctuation, everything to lower case

def process(name):

return "".join(c for c in name if c not in string.punctuation).lower()

processed_node_name = process(node_name)

processed_module_name = process(module_name)

# module names start with 'module.' in distributed data-parallel training, but fx graph node names don't; check for both

distributed_node_name = "module." + node_name

distributed_processed_node_name = "module" + processed_node_name

return (

(node_name == module_name)

or (distributed_node_name == module_name)

or (processed_node_name == processed_module_name)

or (distributed_processed_node_name == processed_module_name)

)

def replicate_sequence(sequence, replications):

"""Replicate a permutation to apply it to an even multiple of channel counts"""

replicated_sequence = []

for rep in range(replications):

offset = len(sequence) * rep

for c in sequence:

replicated_sequence.append(c + offset)

return replicated_sequence

class Permutation:

__model = None

__sparse_parameters = []

__allow_permutation = False

__all_parameters = []

__verbosity = 0 ## 0: errors only, 1: also high-level details, warnings, 2: also intermediate steps, 3: everything

__params_permuted_in_C = []

__params_permuted_in_K = []

__unpermuted_dims = []

__save_permutation_graph = False

__permutation_output_dir = ""

__manual_seed = None

__tcpstore_port = 2341

# these module types may be the target of permutations (have potentially sparse weights or are attributes with no parents)

__permutation_target_module_types = [

"torch.nn.modules.conv.Conv1d",

"torch.nn.modules.conv.Conv2d",

"torch.nn.modules.linear.Linear",

"torch.nn.modules.linear.LazyLinear",

"torch.nn.modules.linear.NonDynamicallyQuantizableLinear",

"torch.nn.modules.activation.MultiheadAttention",

"get_attr",

]

# these module types are not permuted, but must pass any permutation seen by a child's C or passed-thru K to the parents' K

__simple_passthru_module_types = [

"torch.nn.modules.activation.ReLU6",

"torch.nn.modules.activation.ReLU",

"torch.nn.modules.dropout.Dropout",

"torch.nn.modules.dropout.Dropout1d",

"torch.nn.modules.dropout.Dropout2d",

"torch.nn.modules.dropout.Dropout3d",

"torch.nn.modules.dropout.AlphaDropout",

"torch.nn.modules.dropout.FeatureAlphaDropout",

"torch.nn.modules.pooling.MaxPool2d",

"torch.nn.modules.pooling.AdaptiveAvgPool2d",

"torch.nn.modules.pooling.AvgPool2d",

"torch.nn.modules.activation.Hardsigmoid",

"torch.nn.modules.activation.Hardswish",

"torch.nn.modules.activation.GELU",

"torch.nn.modules.normalization.LocalResponseNorm",

"torch.nn.modules.activation.Softmin",

"torch.nn.modules.activation.Softmax",

"torch.nn.modules.activation.Softmax2d",

"torch.nn.modules.activation.LogSoftmax",

"torch.nn.modules.activation.AdaptiveLogSoftmaxWithLoss",

"torch.nn.modules.activation.SiLU",

"torch.nn.modules.activation.Sigmoid",

"concat",

"torch.nn.modules.flatten.Flatten", # if it's a problem, it'll be handled via dimension mismatch check

]

# these module types have parameters that must be permuted along K as well as need to pass the permutation thru to parents' K

__permute_K_and_passthru_module_types = [

"torch.nn.modules.batchnorm.BatchNorm2d",

"torch.nn.modules.normalization.LayerNorm",

"torch.nn.modules.instancenorm.InstanceNorm2d",

"torch.nn.modules.batchnorm.SyncBatchNorm",

]

# these module types cannot be permuted safely (today), and cause neighboring layers to have permutations disabled

__disallow_permutations_module_types = [

"torch.nn.modules.normalization.GroupNorm", # to handle: influence GCD of real children's sibling group

"torch.nn.modules.linear.Bilinear", # need to permute one input along in1_features and the other along in2_features

"torch.nn.modules.activation.GLU", # may work OOTB, but might need to explicitly handle dimsionality change

]

@classmethod

def set_identical_seed(cls, identical_seed=1):

"""Make all GPUs in DDP use the same seed to find identical permutations and not require syncing parameters later"""

if cls.__verbosity > 0:

print(

"[set_identical_seed] Set the identical seed: {:} for all GPUs to make sure the same results generated in permutation search".format(

identical_seed

)

)

cls.__manual_seed = identical_seed

cls.reset_seed()

@classmethod

def reset_seed(cls):

"""To find the same permutations no matter how many GPUs are used, we reset the seed before every search"""

identical_seed = cls.__manual_seed

assert identical_seed is not None, "Must call set_identical_seed() before it can be reset"

torch.manual_seed(identical_seed)

torch.cuda.manual_seed(identical_seed)

import random

np.random.seed(identical_seed)

random.seed(identical_seed)

torch.backends.cudnn.deterministic = True

torch.backends.cudnn.benchmark = False

@classmethod

def set_tcpstore_port(cls, tcpstore_port):

"""Override the default port if it is in use in a distributed training session"""

cls.__tcpstore_port = tcpstore_port

if cls.__verbosity > 0:

print(f"[set_tcpstore_port] TCPStore port set to {cls.__tcpstore_port} .")

@classmethod

def set_permutation_saving_params(

cls,

allow_permutation=False,

save_permutation_graph=False,

permutation_output_dir=".",

):

"""This function is used to set the permutation saving related parameters."""

cls.__allow_permutation = allow_permutation

cls.__save_permutation_graph = save_permutation_graph

cls.__permutation_output_dir = permutation_output_dir

if cls.__verbosity > 0:

print(

f"[permutation_lib][set_permutation_saving_param] Set permutation saving related parameters\n\tAllow permutation: {cls.__alow_permutation}\n\tSave permutation graphs: {cls.__save_permutation_graph}\n\tPermutation graphs saving dir: {cls.__permutation_output_dir}"

)

@classmethod

def set_permutation_params_from_asp(cls, model, sparse_parameters, all_parameters, verbosity):

"""This function is used to set the permutation needed parameters from ASP class."""

cls.__verbosity = verbosity

if cls.__verbosity > 0:

print("[set_permutation_params_from_asp] Set permutation needed parameters")

cls.__model = model

cls.__sparse_parameters = sparse_parameters

cls.__all_parameters = all_parameters

if cls.__verbosity > 1:

sparse_param_names = [

module_name + ":" + p_name

for (

module_name,

module,

p_name,

p,

mask,

pruned,

) in cls.__sparse_parameters

]

all_param_names = [

module_name + ":" + p_name

for (module_name, module, p_name, p) in cls.__all_parameters

]

print(

f"\tSparse parameter names: {sparse_param_names}\n\tAll parameter names: {all_param_names}"

)

cls.__params_permuted_in_C = []

cls.__params_permuted_in_K = []

cls.__unpermuted_dims = []

@classmethod

def permute_model(

cls,

model,

dump_fx_graph=False,

save_dumped_fx_graph="./model_permutation_graph.json",

):

"""Permute a model's weights in order to maintain more magnitude after enforcing the sparsity constraint."""

if cls.__verbosity > 0:

print("\n[permute_model] Permuting the model")

# extract the output_dir, so all the intermediate fx_graph can be saved under that path

extract_output_dir = os.path.split(save_dumped_fx_graph)[0]

cls.__permutation_output_dir = extract_output_dir

fx_graph, success_in_build_fx_graph = cls.build_fx_graph(

model,

dump_fx_graph=dump_fx_graph,

save_dumped_fx_graph=save_dumped_fx_graph,

)

if success_in_build_fx_graph:

fx_graph_after_init_flags = cls.init_permutation_flags(fx_graph)

fx_graph_after_find_real_parents = cls.find_real_parents(fx_graph_after_init_flags)

fx_graph_after_find_real_children = cls.find_real_children(

fx_graph_after_find_real_parents

)

fx_graph_after_making_groups = cls.make_sibling_coparent_groups(

fx_graph_after_find_real_children

)

fx_graph_after_fixup_concats = cls.fixup_concats(fx_graph_after_making_groups)

fx_graph_after_enforce_dimension_agreement = cls.enforce_dimension_agreement(

fx_graph_after_fixup_concats

)

fx_graph_after_propagate_flags = cls.propagate_permutation_flags(

fx_graph_after_enforce_dimension_agreement

)

start_time_search_for_good_permutation = time.perf_counter()

fx_graph_after_find_permutations = cls.find_permutations(fx_graph_after_propagate_flags)

if torch.distributed.is_initialized():

if cls.__verbosity > 0:

duration_search_for_good_permutation = (

time.perf_counter() - start_time_search_for_good_permutation

)

print(

f"[permute_model] Rank {torch.distributed.get_rank()} completed search in {duration_search_for_good_permutation:.2f}s, waiting for others.",

force=True,

)

torch.distributed.barrier()

duration_search_for_good_permutation = (

time.perf_counter() - start_time_search_for_good_permutation

)

if cls.__verbosity > 0:

print(

"\n[permute_model] Take {:.4f} seconds to finish search_for_good_permutation function.".format(

duration_search_for_good_permutation

)

)

fx_graph_after_sync_permutations = cls.sync_permutations(

fx_graph_after_find_permutations

)

fx_graph_after_apply_permutations = cls.apply_permutations(

fx_graph_after_sync_permutations

)

cls.check_graph_for_unpermuted_nodes(fx_graph_after_apply_permutations)

fx_graph = fx_graph_after_apply_permutations

if cls.__save_permutation_graph:

cls.save_graph_to_json(

fx_graph,

save_dumped_graph_path_with_name=os.path.join(

cls.__permutation_output_dir, "./model_graph_permutation_graph.json"

),

) # save the intermediate graph as JSON file for debugging

return success_in_build_fx_graph

@classmethod

def get_permutation_stats(cls):

"""Return statistics for how many permutations were applied in various dimensions, used for testing"""

return (

cls.__params_permuted_in_C,

cls.__params_permuted_in_K,

cls.__unpermuted_dims,

)

@classmethod

def apply_permutation_in_C_dim(cls, node_name, permutation_sequence, dryrun):

"""This function is used to permutation for a node in C dim. (Only need to handle the weight of the node)"""

if cls.__verbosity > 1 and dryrun:

print(

"[apply_permutation_in_C_dim] Permutation for node: '{:}' in C dim".format(

node_name

)

)

if len(permutation_sequence) == 0:

if cls.__verbosity >= 0:

print(

f"ERROR: [apply_permutation_in_C_dim] the permutation sequence for node {node_name} is empty, fail to apply permutation in C dim."

)

return False

is_node_in_sparse_parameters = False

success_permutation = False

for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:

if node_name_matches(node_name, module_name):

if cls.__verbosity > 2 and dryrun:

print(

"[apply_permutation_in_C_dim] find the node: '{:}' '{:}' in cls.__sparse_parameters, succeed to apply permutation in C dim.".format(

node_name, p_name

)

)

is_node_in_sparse_parameters = True

permutation_to_apply = permutation_sequence

if p.shape[1] != len(

permutation_sequence

): # assumed to be grouped convolutions or concatenated weights

if p.shape[1] % len(permutation_sequence) != 0:

return False

permutation_to_apply = replicate_sequence(

permutation_sequence, p.shape[1] // len(permutation_sequence)

)

if not dryrun:

p.data.copy_(p[:, permutation_to_apply, ...])

cls.__params_permuted_in_C.append(node_name + "." + p_name)

success_permutation = True

if not is_node_in_sparse_parameters:

# A special case: if the node itself not in sparse_module_names but one of its real_siblings in sparse_module_names, then the node will not do the permutation search, but it may need to apply the offline permutation in C dim according to the searched permutation sequence from its real_siblings in sparse_module_names

try:

for (

module_name_from_all_parameters,

module_from_all_parameters,

p_name_from_all_parameters,

p_from_all_parameters,

) in cls.__all_parameters:

if (

node_name_matches(node_name, module_name_from_all_parameters)

and p_name_from_all_parameters == "weight"

):

if cls.__verbosity > 3 and dryrun:

print(

"[apply_permutation_in_C_dim] cannot find the node: '{:}' '{:}' in cls.__sparse_parameters, but can find in cls.__all_parameters.".format(

node_name, p_name_from_all_parameters

)

)

permutation_to_apply = permutation_sequence

if p_from_all_parameters.shape[1] != len(

permutation_sequence

): # assumed to be grouped convolutions

if p_from_all_parameters.shpae[1] % len(permutation_sequence) != 0:

return False

permutation_to_apply = replicate_sequence(

permutation_sequence,

p_from_all_parameters.shape[1] // len(permutation_sequence),

)

if not dryrun:

p_from_all_parameters.data.copy_(

p_from_all_parameters[:, permutation_to_apply, ...]

)

cls.__params_permuted_in_C.append(

node_name + "." + p_name_from_all_parameters

)

success_permutation = True

if cls.__verbosity > 2 and dryrun:

print(

"[apply_permutation_in_C_dim] cannot find the node: '{:}' in cls.__sparse_parameters, after trying with cls.__all_parameters, succeed to apply permutation in C dim.".format(

node_name

)

)

except:

success_permutation = False

if cls.__verbosity >= 0:

print(

"ERROR: [apply_permutation_in_C_dim] cannot find the node: '{:}' in cls.__sparse_parameters, after trying with cls.__all_parameters, still fail to apply permutation in C dim.".format(

node_name

)

)

return success_permutation

@classmethod

def permute_attr(cls, node_name, permutation_sequence, fx_graph, dryrun):

"""Permute a node's attributes. Somewhat hacky, assumes that we'll find exactly one dimension with a length matching the permutation's"""

assert "attr" in fx_graph[node_name].keys()

attr = fx_graph[node_name]["attr"]

if cls.__verbosity > 1:

print(f"Found attribute {node_name} of shape {attr.shape}")

found_perm = False

for dim in range(len(attr.shape)):

if attr.shape[dim] == len(permutation_sequence):

if found_perm:

if cls.__verbosity > 0:

print(

f"\tWARNING: {node_name} has already been permuted, but it's trying to happen again along another dimension {dim}."

)

return False

found_perm = True

if cls.__verbosity > 1 and dryrun:

print(f"\tpermuting along dimension {dim}")

if not dryrun:

# permute the dimension of interest to the front, permute within that dimension, then reset it

order = [c for c in range(len(attr.shape))]

order[0] = dim

order[dim] = 0

prmt = tuple(order)

temp_weight = torch.clone(attr)

temp_weight = torch.permute(temp_weight, prmt)

temp_weight.copy_(temp_weight[permutation_sequence, ...])

temp_weight = torch.permute(temp_weight, prmt)

attr.data.copy_(temp_weight)

cls.__params_permuted_in_K.append(node_name + "_" + str(dim))

return found_perm

@classmethod

def apply_permutation_in_K_dim(cls, node_name, permutation_sequence, fx_graph, dryrun):

"""This function is used to permutation for a node in K dim. (Need to handle the weight/bias/running_mean/running_var of the node)"""

if cls.__verbosity > 1:

print(

"[apply_permutation_in_K_dim] Permutation for node: '{:}' in K dim".format(

node_name

)

)

if len(permutation_sequence) == 0:

if cls.__verbosity >= 0:

print(

"ERROR: [apply_permutation_in_K_dim] the permutation sequence is empty, fail to apply permutation in K dim."

)

return False

# permute attribute nodes

if "attr" in fx_graph[node_name].keys():

return cls.permute_attr(node_name, permutation_sequence, fx_graph, dryrun)

# if we didn't store the attribute already, look in the modules' parameters

is_node_in_all_parameters = False

success_permutation = False

for module_name, module, p_name, p in cls.__all_parameters:

if node_name_matches(node_name, module_name):

if cls.__verbosity > 1 and dryrun:

print(

"[apply_permutation_in_K_dim] find the node: '{:}' with '{:}' in cls.__all_parameters, may succeed to apply permutation in K dim.".format(

node_name, p_name

)

)

is_node_in_all_parameters = True

permutation_to_apply = permutation_sequence

if p.shape[0] != len(permutation_sequence): # assumed to be grouped convolutions

if cls.__verbosity > 2 and dryrun:

print(

f"Mismatch in K dimension between found module {module_name} {p_name} for node {node_name}: permutation length {len(permutation_sequence)} but parameter shape in K {p.shape[0]}"

)

if p.shape[0] % len(permutation_sequence) != 0:

return False

permutation_to_apply = replicate_sequence(

permutation_sequence, p.shape[0] // len(permutation_sequence)

)

if cls.__verbosity > 1 and dryrun:

print(

"[apply_permutation_in_K_dim] the node: '{:}' with shape: '{:}' required replicating the permutation sequence with len '{:}' {:} times to succeed in applying the permutation in the K dimension.".format(

node_name,

p.shape,

len(permutation_sequence),

p.shape[0] // len(permutation_sequence),

)

)

else:

if cls.__verbosity > 1 and dryrun:

print(

"[apply_permutation_in_K_dim] the node: '{:}' with shape: '{:}', can match the size of permutation sequence with len: '{:}', succeed to apply permutation in K dim.".format(

node_name, p.shape, len(permutation_sequence)

)

)

if not dryrun:

p.data.copy_(p[permutation_to_apply, ...])

cls.__params_permuted_in_K.append(node_name + "." + p_name)

success_permutation = True

if not is_node_in_all_parameters:

if cls.__verbosity >= 0:

print(

"ERROR: [apply_permutation_in _K_dim] cannot find the node: '{:}' in cls.__all_parameters, fail to apply permutation in K dim.".format(

node_name

)

)

success_permutation = False

return success_permutation

@classmethod

def check_graph_for_unpermuted_nodes(cls, fx_graph):

"""Make sure that all permutable nodes/parameters were actually permuted and all GPUs agree"""

for node_name in fx_graph.keys():

node = fx_graph[node_name]

if "C_permutable" in node.keys() and node["C_permutable"] and not node["C_permuted"]:

sibling_group_id = node["sibling_group_id"]

if (

node["is_real"]

and cls.__group_data["skipped_sibling_groups"][sibling_group_id] is None

):

if cls.__verbosity >= 0:

print(

f"{node_name} was C_permutable in a not skipped sibling group but was not permuted along C! {node}"

)

cls.__unpermuted_dims.append(node_name + "_C")

if "K_permutable" in node.keys() and node["K_permutable"] and not node["K_permuted"]:

coparent_group_id = node["coparent_group_id"]

if (

node["is_real"]

and cls.__group_data["skipped_coparent_groups"][coparent_group_id] is None

):

if cls.__verbosity >= 0:

print(

f"{node_name} was K_permutable in a not skipped coparent group but was not permuted along K! {node}"

)

cls.__unpermuted_dims.append(node_name + "_K")

if cls.__verbosity > 0:

print(

f"[check_graph_for_unpermuted_nodes] found nodes that missed permutations along {len(cls.__unpermuted_dims)} dimensions."

)

# make sure all GPUs agree

if torch.distributed.is_initialized():

cls.__unpermuted_dims = sorted(cls.__unpermuted_dims)

rank = torch.distributed.get_rank()

world_size = torch.distributed.get_world_size()

dist_store = torch.distributed.TCPStore(

"127.0.0.1", cls.__tcpstore_port, world_size, rank == 0

)

torch.distributed.barrier()

dist_store.set(str(rank), ",".join(cls.__unpermuted_dims))

torch.distributed.barrier()

if rank == 0:

my_list = dist_store.get("0").decode()

for peer in range(1, world_size):

peer_list = dist_store.get(str(peer)).decode()

assert my_list == peer_list, (

f"peer {peer} disagreed with rank 0's list of unpermuted nodes: \n{my_list}\n{peer_list}"

)

@classmethod

def find_sparse_parameters_for_node(cls, node_name):

"""If the node has parameters that are in the trackd sparse parameter list, find them and reshape to a 2D tensor with channels last"""

node_weight = None

# check the sparse parameters

for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:

if node_name_matches(node_name, module_name):

node_weight = torch.zeros_like(p)

node_weight.copy_(p)

# if we found something, reshape to concatenate along the same dimension

if node_weight is not None:

# Need to handle the concat for layers with different R & S

shape = node_weight.shape

# 1d-tensor

if len(shape) == 1:

node_weight = node_weight.view(1, shape[0])

# 2d-tensor (K, C)

elif len(shape) == 2:

node_weight = node_weight.view(shape[0], shape[1])

# 3d-tensor (K, C, R)

elif len(shape) == 3:

node_weight = (

node_weight.permute(0, 2, 1).contiguous().view(shape[0] * shape[2], shape[1])

)

# 4d-tensor (K, C, R, S)

elif len(shape) == 4:

# convs

node_weight = (

node_weight.permute(2, 3, 0, 1)

.contiguous()

.view(shape[2] * shape[3] * shape[0], shape[1])

)

return node_weight

@classmethod

def find_permutation_for_matrix_group(cls, matrix_group):

"""Find a good permutation for some matrix (which may be concatenated matrices that require the same permutation)"""

if cls.__verbosity > 1:

print(

f"Searching for a good permutation for this sibling group of shape {matrix_group.shape}"

)

permutation_found = False

num_channels = matrix_group.shape[1]

group_permutation = [c for c in range(num_channels)]

# automatic check for skipping the permutation search process

original_magnitude = (torch.abs(matrix_group)).sum(dtype=torch.float64)

pruned_magnitude = sum_after_2_to_4(matrix_group.cpu().detach().numpy())

diff_ratio = abs(original_magnitude - pruned_magnitude) / original_magnitude

epsilon = 1e-3

if cls.__verbosity > 1:

print(

"\n[search_for_good_permutation] Original element abs sum: {:}, Pruned element abs sum: {:}, Diff ratio: {:}".format(

original_magnitude, pruned_magnitude, diff_ratio

)

)

start_time_accelerated_search_for_good_permutation = time.perf_counter()

if diff_ratio < epsilon:

if cls.__verbosity > 2:

print(

"[search_for_good_permutation] Original element abs sum is almost same as the pruned element abs sum, further permutation search will not help, skipping!"

)

else:

if cls.__verbosity > 2:

print(

"[search_for_good_permutation] Original element abs sum is different from the pruned element abs sum, further permutation search will help, continue with the permutation search!"

)

# call the permutation search CUDA kernels as ASP extension.

# users can provide prefer search strategy by providing a valid 'search_options' as a dictionary,

# or users can implement their customized 'accelerated_search_for_good_permutation' function.

search_options = {}

# No.1 Strategy: Exhaustive Search

search_options["strategy"] = "exhaustive"

search_options["stripe_group_size"] = 8

search_options["escape_attempts"] = 100

# No.2 Strategy: Progressive Channel Swap Search

# search_options['strategy'] = 'progressive channel swap'

# search_options['progressive_search_time_limit'] = 10

# search_options['improvement_threshold'] = 1e-9

# permutation search time is too long for matrix_group with large channel num

# change from Exhaustive Search to Progressive Channel Swap Search based on input matrix_group size

if num_channels > 2048:

search_options = {}

search_options["strategy"] = "progressive channel swap"

search_options["progressive_search_time_limit"] = 120

search_options["improvement_threshold"] = 1e-9

if cls.__verbosity > 1:

print(f"[search_for_good_permutation] search options: {search_options}")

group_permutation = accelerated_search_for_good_permutation(

matrix_group, options=search_options, verbosity=cls.__verbosity

)

permutation_found = True

if cls.__verbosity > 1:

duration_accelerated_search_for_good_permutation = (

time.perf_counter() - start_time_accelerated_search_for_good_permutation

)

permuted_magnitude = sum_after_2_to_4(

matrix_group.cpu().detach().numpy()[:, group_permutation]

)

print(

"[search_for_good_permutation] Take {:.4f} seconds to finish accelerated_search_for_good_permutation function and with final magnitude {:}.".format(

duration_accelerated_search_for_good_permutation, permuted_magnitude

)

)

return group_permutation, permutation_found

@classmethod

def skip_sibling_group(cls, fx_graph, sibling_group_id, reason):

"""Keep track of sibling groups that do not have permutations applied"""

# grab a parent to get the coparent group id

sibling_group = cls.__group_data["sibling_groups"][sibling_group_id]

a_sibling = list(sibling_group)[0]

a_parent = fx_graph[a_sibling]["real_parents"][0]

coparent_group_id = fx_graph[a_parent]["coparent_group_id"]

if cls.__verbosity > 1:

print(

f"Skipping permutations for Sibling Group {sibling_group_id} and Coparent Group {coparent_group_id}: {reason}"

)

cls.__group_data["skipped_sibling_groups"][sibling_group_id] = reason

cls.__group_data["skipped_coparent_groups"][coparent_group_id] = reason

@classmethod

def collect_sparse_weights(cls, fx_graph, sibling_group, sibling_group_C_param):

"""Gather all sparse weights for a sibling group (to serve as input to the permutation search)"""

matrix_group = None

for sibling in sibling_group:

node_weight = cls.find_sparse_parameters_for_node(sibling)

if node_weight is not None:

# reshape due to siblings with grouped convolutions of different sizes

assert node_weight.shape[1] % sibling_group_C_param == 0, (

f"sibling {sibling}'s weights' C={node_weight.shape[1]} must be even multiple of the sibling group's C parameter {sibling_group_C_param}"

)

node_weight = torch.reshape(node_weight, (-1, sibling_group_C_param))

if matrix_group is None:

matrix_group = node_weight

else:

try:

matrix_group = torch.cat(

(matrix_group, node_weight), dim=0

) # concat the weights in the K dimension, keep the same C dimension

except:

if cls.__verbosity >= 0:

print(

"ERROR: [search_for_good_permutation][warning] cannot merge the weight for node: '{:}', with its weight shape: '{:}', the matrix_group shape: '{:}'.".format(

sibling, node_weight.size(), matrix_group.size()

)

)

continue

if cls.__verbosity > 2:

print(

"[search_for_good_permutation] have merged the weight for node: '{:}', with its weight shape: '{:}', the matrix_group shape: '{:}'.".format(

sibling, node_weight.size(), matrix_group.size()

)

)

else:

if cls.__verbosity > 2:

print(

f"[search_for_good_permutation] not adding dense weights for node {sibling} to the group"

)

return matrix_group

@classmethod

def find_sibling_group_permutation(cls, fx_graph, sibling_group_id):

""" "Find a good permutation for some sibling group"""

if cls.__verbosity > 1:

print(f"Finding permutation for sibling group {sibling_group_id}")

cls.reset_seed()

sibling_group = cls.__group_data["sibling_groups"][sibling_group_id]

sibling_group_C_param = int(cls.__group_data["sibling_group_C_params"][sibling_group_id])

if sibling_group_C_param % 4 != 0 or sibling_group_C_param < 8:

cls.skip_sibling_group(

fx_graph, sibling_group_id, f"Useless C: {sibling_group_C_param}"

)

return

# collect *sparse* weights from all siblings, get the coparent group

matrix_group = cls.collect_sparse_weights(fx_graph, sibling_group, sibling_group_C_param)

# early-out if no siblings are sparse

if matrix_group is None:

cls.skip_sibling_group(fx_graph, sibling_group_id, "Dense")

return

# find a good permutation

group_permutation, found = cls.find_permutation_for_matrix_group(matrix_group)

# if no permutation was found, we didn't need it (input already sparse)

if not found:

cls.skip_sibling_group(fx_graph, sibling_group_id, "Not needed")

return

if cls.__verbosity > 2:

print(f"Permutation for sibling group {sibling_group_id}: {group_permutation}")

cls.__group_data["sibling_group_permutations"][sibling_group_id] = group_permutation

@classmethod

def permute_sibling_group(cls, fx_graph, sibling_group_id, group_permutation):

"""Apply a permutation to some sibling group"""

if cls.__verbosity > 1:

print(f"Attempting to permute sibling group {sibling_group_id}")

sibling_group = cls.__group_data["sibling_groups"][sibling_group_id]

# apply the permutation in two steps: first, a dry run to find any issues.

# if there were no issues, actually apply the permutation in the second step.

success = True

coparent_group_id = None

for dryrun in [True, False]:

# apply that permutation to the siblings' C dimension

for sibling in sibling_group:

assert fx_graph[sibling]["C_permutable"] and not fx_graph[sibling]["C_permuted"]

sibling_permuted = cls.apply_permutation_in_C_dim(

sibling, group_permutation, dryrun

)

if dryrun:

success = success and sibling_permuted

else:

assert sibling_permuted, "shouldn't fail permuting siblings after the dry run"

fx_graph[sibling]["C_permuted"] = sibling_permuted

a_parent = fx_graph[sibling]["real_parents"][0]

if coparent_group_id is None:

coparent_group_id = fx_graph[a_parent]["coparent_group_id"]

else:

assert coparent_group_id == fx_graph[a_parent]["coparent_group_id"], (

f"parent {a_parent} must belong to the same coparent group {coparent_group_id}, not {fx_graph[a_parent]['coparent_group_id']}"

)

# grab the parents (and co-parents) and apply to their K dimension

coparents = cls.__group_data["coparent_groups"][coparent_group_id]

for coparent in coparents:

assert fx_graph[coparent]["K_permutable"] and not fx_graph[coparent]["K_permuted"]

coparent_permuted = cls.apply_permutation_in_K_dim(

coparent, group_permutation, fx_graph, dryrun

)

if dryrun:

success = success and coparent_permuted

else:

assert coparent_permuted, "shouldn't fail permuting coparents after the dry run"

fx_graph[coparent]["K_permuted"] = coparent_permuted

children_permuted = cls.apply_permutation_in_K_dim_to_children(

fx_graph, coparent, group_permutation, dryrun

)

if dryrun:

success = success and children_permuted

else:

assert children_permuted, (

"shouldn't fail permuting coparents' children after the dry run"

)

if not success:

cls.skip_sibling_group(fx_graph, sibling_group_id, "dryrun_failure")

if cls.__verbosity > 0:

print(

f"There was an issue permuting sibling group {sibling_group_id}, skipping it to preserve network quality."

)

break

@classmethod

def apply_permutation_in_K_dim_to_children(cls, fx_graph, node_name, permutation, dryrun):

"""Apply a permutation along K to the children of some node"""

success = True

children = fx_graph[node_name]["children"]

if cls.__verbosity > 2 and dryrun:

print(f"Applying a permutation in K to children of {node_name} : {children}")

# apply the permutation along K to children as necessary

for child in children:

if "is_real" in fx_graph[child].keys() and fx_graph[child]["is_real"]:

if cls.__verbosity > 3 and dryrun:

print(f"\tFound a real child {child}, not permuting it or its children along K")

else:

if (

"module_type" not in fx_graph[child].keys()

or fx_graph[child]["module_type"] == "None"

):

if cls.__verbosity > 3 and dryrun:

print(f"\tPermuting children of non-module {child} along K")

success = success and cls.apply_permutation_in_K_dim_to_children(

fx_graph, child, permutation, dryrun

)

elif not fx_graph[child]["C_permutable"]:

if fx_graph[child]["K_permutable"] and not fx_graph[child]["K_permuted"]:

if cls.__verbosity > 2 and dryrun:

print(f"\tPermuting {child} along K")

child_permuted = cls.apply_permutation_in_K_dim(

child, permutation, fx_graph, dryrun

)

success = success and child_permuted

if not dryrun:

fx_graph[child]["K_permuted"] = child_permuted

assert fx_graph[child]["K_passthru"]

if fx_graph[child]["K_passthru"]:

success = success and cls.apply_permutation_in_K_dim_to_children(

fx_graph, child, permutation, dryrun

)

else:

if cls.__verbosity >= 0:

print(

f"\t!! ERROR {child} was a not real module that was not K_passthru"

)

return success

@classmethod

def defer_prints(cls):

"""Collect prints from this rank in distributed mode to avoid interleaved output"""

if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:

cls.__new_stdout = io.StringIO(str(torch.distributed.get_rank()))

cls.__builtin_print = __builtin__.print

def deferred_print(*args, **kwargs):

try: # see if torchvision examples has suppressed other ranks with the force argument

cls.__builtin_print(*args, file=cls.__new_stdout, force=True, **kwargs)

except:

cls.__builtin_print(*args, file=cls.__new_stdout, **kwargs)

__builtin__.print = deferred_print

@classmethod

def resume_prints(cls):

"""Emit the collected outputs from this rank, resume immediate printing"""

if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:

output = cls.__new_stdout.getvalue()

__builtin__.print = cls.__builtin_print

try:

print(output, force=True)

except:

print(output)

@classmethod

def find_permutations(cls, fx_graph):

"""Search for permutations for all sibling groups"""

for sibling_group_id in cls.__group_data["sibling_groups"].keys():

search_this_group = True

if torch.distributed.is_initialized():

rank = torch.distributed.get_rank()

world_size = torch.distributed.get_world_size()

if sibling_group_id % world_size != rank: