| #------------------------------------------------------------- |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| #------------------------------------------------------------- |
| # Independent Subnet Training (IST) |
| # |
| # This builtin implements independent subnet training as a |
| # second-order function. It orchestrates distributed / parallel |
| # training over disjoint subnets using parfor, while delegating |
| # architecture-specific logic to user-provided functions. |
| # ------------------------------------------------------------ |
| # INPUT: |
| # model : initial model parameters. A list of matrices for NN. |
| # features : X |
| # labels : Y |
| # val_features : validation X |
| # val_labels : validation Y |
| # upd : computes gradients and performs optimizer step |
| # agg : aggregation logic to combine updates of shared parameters (across subnets) |
| # epochs : number of epochs |
| # batchsize : batchsize for training |
| # j : number of gradient steps until aggregation -> determines length of the IST round (aggregation frequency) |
| # numSubnets : number of independent subnets/workers |
| # hyperparams : list of hyperparameters (e.g. lr, reg, mask params, etc.) |
| # verbose : print progress (boolean) |
| # paramsPerLayer : amount of parameters each layer consists of |
| # fullyConnectedLayers : list of all FC layer indices (starting at idx=1) |
| # |
| # OUTPUT: |
| # model_out : trained model parameters (IST: W) |
| # |
| # ASSUMPTION: |
| # - the last layer is the output layer |
| # ------------------------------------------------------------ |
| |
| m_independentSubnetTrain = function( |
| list[unknown] model, |
| matrix[double] features, |
| matrix[double] labels, |
| matrix[double] val_features, |
| matrix[double] val_labels, |
| string upd, |
| string agg, |
| int epochs, |
| int batchsize, |
| int j, |
| int numSubnets, |
| list[unknown] hyperparams, |
| boolean verbose, |
| int paramsPerLayer, |
| list[int] fullyConnectedLayers |
| ) |
| return (list[unknown] trained_model) |
| { |
| # ------------------------------------------------------------ |
| # Setup |
| # ------------------------------------------------------------ |
| model_out = model |
| |
| P = length(model) |
| N = nrow(features) |
| if (P %% paramsPerLayer != 0) stop("Model length not divisible by paramsPerLayer") |
| L = as.integer(P / paramsPerLayer) # total layers |
| |
| # I. determine shared parameters |
| isSharedParam = matrix(0, 1, P) |
| |
| # - create mask for all FC layers |
| fcLayers = fullyConnectedLayers |
| isFC = matrix(0, rows=1, cols=L) |
| for (i in 1:length(fcLayers)) { |
| idx = as.integer(as.scalar(fcLayers[i])) |
| isFC[1, idx] = 1 # TODO vectorize |
| } |
| |
| # - expand layer mask across all parameters |
| isFC_rep = isFC |
| for (r in 2:paramsPerLayer) { |
| isFC_rep = cbind(isFC_rep, isFC) |
| } |
| if (ncol(isFC_rep)!=P) stop("Dimension mismatch for FC layer mask.") |
| |
| # - all non-FC layers are shared |
| isSharedParam = 1 - isFC_rep |
| |
| # - edge case: FC bias parameters are shared in: output layer or at the end of a FC block |
| for (paramId in seq(2, paramsPerLayer, 2)) { # iterate bias blocks only |
| for (l in 1:L) { |
| if (as.scalar(isFC[1,l])==1 & l==L) { |
| p_out_bias = (paramId - 1) * L + L # output bias is shared across subnets |
| isSharedParam[1, p_out_bias] = 1 # TODO vectorize |
| } |
| else if (as.scalar(isFC[1,l])==1 & l<L & as.scalar(isFC[1,l+1])==0) { |
| p_out_bias = (paramId - 1) * L + l # end of FC block's bias is shared across subnets |
| isSharedParam[1, p_out_bias] = 1 # TODO vectorize |
| } |
| } |
| } |
| if (ncol(isSharedParam) != P) stop("isSharedParam dimension mismatch!") |
| |
| # II. calculate update-steps per epoch |
| if (batchsize<=0 | batchsize>N) { |
| stop("Batch size is out of bounds!") |
| } else { |
| stepsPerEpoch = ceil(N / batchsize) |
| } |
| |
| # III. training loop |
| for (epoch in 1:epochs) { |
| if (verbose) print("Entered epoch: " + epoch) |
| |
| # A.) reshuffle indices each epoch |
| allSampleIndicesRandom = order(target=rand(rows=N, cols=1), by=1, decreasing=FALSE, index.return=TRUE) |
| batchIndices = allSampleIndicesRandom[, 1] |
| b = nrow(batchIndices) |
| I = seq(1, b, 1) |
| V = matrix(1, rows=b, cols=1) |
| S = table(I, batchIndices, V, b, N) |
| |
| features_shuffled = S %*% features |
| labels_shuffled = S %*% labels |
| |
| # B.) iterate IST rounds |
| for (step in seq(1, stepsPerEpoch, j)) { |
| if (verbose) print("Starting new IST round at step: " + step) |
| round_model = model_out # prevent accidental mutation of model_out |
| |
| # 1.) create masks for all subnets |
| [masks, masks_meta_info] = ist_create_disjoint_masks(round_model, numSubnets, L, fcLayers, paramsPerLayer, isFC, verbose) |
| |
| # 2.) preallocate list to store all subnets TODO move outside epoch loop? to prevent constantly allocating... |
| updatedSubnets = list() |
| updatedSubnetsMasks = list() |
| for (s in 1:numSubnets) { |
| updatedSubnets = append(updatedSubnets, list()) |
| updatedSubnetsMasks = append(updatedSubnetsMasks, list()) |
| } |
| |
| # 3) create a template for each subnet based on input model (allows indexing in subsequent parfor-loop) TODO move outside epoch loop? to prevent constantly allocating... |
| subnetModelTemplate = list() |
| subnetModelMaskTemplate = list() |
| for (pIdx in 1:P) { |
| subnetModelTemplate = append(subnetModelTemplate, as.matrix(model[pIdx])) |
| subnetModelMaskTemplate = append(subnetModelMaskTemplate, as.matrix(model[pIdx])) |
| } |
| |
| # local optimization steps / IST round |
| localSteps = min(j, (stepsPerEpoch-step+1)) |
| |
| # 4.) obtain all minibatches for this IST round (doing it once prevents parfor confusion) |
| shuffled_features = list() |
| shuffled_labels = list() |
| for (localStep in 1:localSteps) { |
| mb = (step-1) + localStep |
| mb_local = mb-1 |
| start = mb_local*batchsize + 1 |
| end = min(mb*batchsize, N) |
| |
| Xb = features_shuffled[start:end, 1:ncol(features_shuffled)] |
| yb = labels_shuffled[start:end, 1:ncol(labels_shuffled)] |
| shuffled_features = append(shuffled_features, Xb) |
| shuffled_labels = append(shuffled_labels, yb) |
| } |
| |
| # 5.) perform 'j' local gradient steps for each subnet |
| parfor (subnet in 1:numSubnets) { |
| |
| # a.) obtain masked subnet |
| subnet_model = subnetModelTemplate |
| subnet_model_mask = subnetModelMaskTemplate |
| for (subnet_p in 1:length(round_model)) { |
| param_start_idx = as.integer(as.scalar(masks_meta_info[subnet_p,1])) |
| param_end_idx = as.integer(as.scalar(masks_meta_info[subnet_p,2])) |
| param_rows = as.integer(as.scalar(masks_meta_info[subnet_p,3])) |
| param_cols = as.integer(as.scalar(masks_meta_info[subnet_p,4])) |
| |
| vec = masks[subnet, param_start_idx:param_end_idx] |
| param_mask = matrix(vec, rows=param_rows, cols=param_cols, byrow=TRUE) |
| param = as.matrix(round_model[subnet_p]) |
| |
| subnet_model[subnet_p] = list(param * param_mask) # TODO sparse masking! dense masking will probably increase computational efficiency |
| subnet_model_mask[subnet_p] = list(param_mask) |
| } |
| |
| # b.) local optimization steps / IST round |
| for (localStep in 1:localSteps) { |
| feat = as.matrix(shuffled_features[localStep]) |
| lab = as.matrix(shuffled_labels[localStep]) |
| |
| # compute gradients for subnet s + apply update step on owned params (only) |
| subnet_model = as.list(evalList(upd, list(model=subnet_model, mask=subnet_model_mask, features=feat, labels=lab, hyperparams=hyperparams))) |
| } |
| |
| # c.) save updated subnet and mask |
| updatedSubnets[subnet] = list(subnet_model) |
| updatedSubnetsMasks[subnet] = list(subnet_model_mask) |
| } |
| if (verbose) print("All subnets have run successfully.") |
| |
| # 6.) aggregate updates into global model (i.e. model_out) |
| for (p in 1:P) { |
| if (as.scalar(isSharedParam[1, p])==1) { |
| # construct full model update by aggregating shared parameter updates from all subnets |
| subnetParams = list() |
| subnetMasks = list() |
| for (s in 1:numSubnets) { |
| subnet = as.list(updatedSubnets[s]) |
| subnetMask = as.list(updatedSubnetsMasks[s]) |
| |
| subnetParams = append(subnetParams, as.matrix(subnet[p])) |
| subnetMasks = append(subnetMasks, as.matrix(subnetMask[p])) |
| } |
| |
| # aggregate shared parameters based on provided function |
| averagedUpdatedParam = eval(agg, list(initialParam=as.matrix(round_model[p]), allSubnetsParam=subnetParams, allSubnetsMasks=subnetMasks)) |
| round_model[p] = averagedUpdatedParam |
| } |
| else { |
| # construct full model update by filling with disjointly partitioned parameter updates from all subnets |
| initialParam = as.matrix(round_model[p]) |
| updatedParam = matrix(0, nrow(initialParam), ncol(initialParam)) |
| owned = matrix(0, nrow(initialParam), ncol(initialParam)) |
| |
| for (s in 1:numSubnets) { |
| subnet = as.list(updatedSubnets[s]) |
| subnetMask = as.list(updatedSubnetsMasks[s]) |
| |
| owned = owned + as.matrix(subnetMask[p]) |
| updatedParam = updatedParam + as.matrix(subnet[p]) |
| } |
| |
| # SANITY CHECK: |
| max_freq = max(owned) |
| if (max_freq > 1) stop("Overlap detected") |
| |
| round_model[p] = updatedParam |
| } |
| } |
| if (verbose) print("Aggregation of subnets finished. IST round has been successfully executed!") |
| |
| # 7.) update global model (end of the IST round) |
| model_out = round_model |
| } |
| # TODO (potentially): add validation for early stopping etc. |
| } |
| trained_model = model_out |
| } |
| |
| |
| # ---------------------------------------------------------------------------------------------------------------------- |
| # Independent Subnet Masking |
| # |
| # This helper function creates two matrices: one contains all flattened masks, the other contains the info on |
| # how to reconstruct the mask matrices. Each mask is a binary vector indicating which parameters belong to that subnet. |
| # ---------------------------------------------------------------------------------------------------------------------- |
| # INPUT: |
| # model : list of parameter tensors grouped by parameter type i.e. blocks |
| # numSubnets : number of independent subnets/workers (K) |
| # L : total number of layers INCLUDING the output layer (layer indices are assumed to be 1..L) |
| # fullyConnectedLayers : list of all FC layer indices (starting at idx=1) |
| # paramsPerFCLayer : number of parameters / neurons to be partitioned per FC layer |
| # isFC : indicator matrix encoding which layers are FC => isFC[l] ∈ {0,1} |
| # verbose : print progress (boolean) |
| # |
| # OUTPUT: |
| # masks_new : mask matrix defining disjoint neuron ownership across subnets |
| # masks_new_meta : metadata matrix describing the mask layout and ownership mapping |
| # |
| # ASSUMPTIONS: |
| # - neuron ownership is defined via bias vectors |
| # - model is a list of parameter tensors |
| # - trainable parameters are grouped by parameter type i.e. param blocks like (W_l1, W_l2, ..., b_l1, b_l2, ...) |
| # - assumes W and b are always the first two param blocks |
| # - the pattern of optional optimizer state tensors (e.g., vW_l, vb_l) follow the same grouping and always W followed by b |
| # - (output layer & end of FC block) biases are shared -> gradients collide; must be handled by aggregation logic |
| # ---------------------------------------------------------------------------------------------------------------------- |
| |
| ist_create_disjoint_masks = function( |
| list[unknown] model, |
| int numSubnets, |
| int L, |
| list[int] fullyConnectedLayers, |
| int paramsPerFCLayer, |
| Matrix[Double] isFC, |
| boolean verbose |
| ) |
| return ( |
| Matrix[Double] masks_new, |
| Matrix[Double] masks_new_meta |
| ) |
| { |
| P = length(model) |
| |
| # SANITY CHECKS: ensure provided model can be masked correctly |
| if (as.integer(P / paramsPerFCLayer) != L) { |
| stop("Layer/parameter mismatch. Please make sure each layer has the same amount of parameters.") |
| }; |
| if (paramsPerFCLayer < 2 | paramsPerFCLayer %% 2 != 0) { |
| stop("At least 1 pair of W and b needs to be present, as well as parameters need to be W&b pairs.") |
| } |
| |
| # I.) initialize and preallocate masks |
| masks_new_meta = matrix(0, rows=length(model), cols=4) # columns=[start,end,rows,cols] |
| current_position = 1 |
| for (p in 1:length(model)) { |
| M = as.matrix(model[p]) |
| param_length = ncol(M) * nrow(M) # as.scalar(ncol(M)) * as.scalar(nrow(M)) |
| |
| masks_new_meta[p,1] = current_position |
| masks_new_meta[p,2] = current_position + param_length -1 |
| masks_new_meta[p,3] = nrow(M) |
| masks_new_meta[p,4] = ncol(M) |
| |
| current_position = current_position + param_length |
| } |
| mask_size = current_position-1 |
| masks_new = matrix(0, rows=numSubnets, cols=mask_size) # all subnets in one matrix |
| |
| # II.) iterate all layers |
| for (l in 1:L) { |
| |
| # FC layer: create #{numSubnets} disjoint partitions for this layer across all parameters |
| if (as.scalar(isFC[1,l]) == 1) { |
| W = as.matrix(model[l]) |
| b = as.matrix(model[l+L]) |
| H = ncol(W); # bias neurons in layer l |
| |
| # SANITY CHECKS: |
| if (nrow(b) != 1 | ncol(b) != H) { |
| if (verbose) print("Bias shape mismatch!") |
| if (verbose) print("b:", nrow(b), "x", ncol(b)) |
| if (verbose) print("expected: 1 x", H) |
| stop("Invalid bias shape") |
| } |
| if (l!=L & numSubnets>ncol(b)) { # TODO change to next layer is non-FC logic |
| if (verbose) print("More subnets than available neurons in layer:") |
| if (verbose) print(l) |
| stop("Please use a wider model or decrease the amount of subnets.") |
| } |
| |
| # A.) shuffle all neuron indices |
| allNeuronIndicesRandom = order(target=rand(rows=H, cols=1), by=1, decreasing=FALSE, index.return=TRUE) |
| |
| # B.) determine neuron ownership |
| chunk_size = floor(H/numSubnets) |
| remaining_neurons = H - chunk_size * numSubnets |
| amount_active_neurons = matrix(chunk_size, rows=numSubnets, cols=1) |
| if (remaining_neurons > 0) { |
| randomSubnetIndices = order(target=rand(rows=numSubnets, cols=1, seed=-1), by=1, decreasing=FALSE, index.return=TRUE) # TODO replace seed for experiments |
| for (i in 1:remaining_neurons) { |
| sid = as.integer(as.scalar(randomSubnetIndices[i,1])) |
| amount_active_neurons[sid,1] = as.scalar(amount_active_neurons[sid,1]) + 1 # TODO VECTORIZE |
| } |
| } |
| neuron_end_indices = cumsum(amount_active_neurons) |
| neuron_start_indices = neuron_end_indices - amount_active_neurons + 1 |
| |
| # C.) obtain masks for all subnets |
| for(s in 1:numSubnets) { |
| |
| # 1. obtain owned neurons for this layer |
| start = as.integer(as.scalar(neuron_start_indices[s,1])) |
| end = as.integer(as.scalar(neuron_end_indices[s,1])) |
| current_b_indices = allNeuronIndicesRandom[start:end, 1] |
| |
| # 2. create masked bias |
| if(l==L) { # output layer |
| masked_b = matrix(1, rows=1, cols=ncol(b)) |
| } |
| else if (l<L & as.scalar(isFC[1, l+1]) == 0) { # next layer is not FC |
| masked_b = matrix(1, rows=1, cols=ncol(b)) |
| } |
| else { |
| masked_b = matrix(0, rows=1, cols=ncol(b)) |
| for (i in 1:nrow(current_b_indices)) { # TODO VECTORIZE |
| idx = as.integer(as.scalar(current_b_indices[i,1])) |
| masked_b[1, idx] = 1 |
| } |
| } |
| |
| # 3. create masked weight |
| masked_W = matrix(0, rows=nrow(W), cols=ncol(W)) |
| if(l==1) { |
| for (i in 1:nrow(current_b_indices)) { # TODO VECTORIZE |
| idx = as.integer(as.scalar(current_b_indices[i,1])) |
| masked_W[1:nrow(W), idx] = matrix(1, rows=nrow(W), cols=1) |
| } |
| } |
| else if (l>1 & as.scalar(isFC[1, l-1])==0) { # previous layer is not FC |
| for (i in 1:nrow(current_b_indices)) { # TODO VECTORIZE |
| idx = as.integer(as.scalar(current_b_indices[i,1])) |
| masked_W[1:nrow(W), idx] = matrix(1, rows=nrow(W), cols=1) |
| } |
| } |
| else { |
| # obtain active neurons of previous layer |
| p = L + (l-1) |
| start = as.integer(as.scalar(masks_new_meta[p,1])) |
| end = as.integer(as.scalar(masks_new_meta[p,2])) |
| r = as.integer(as.scalar(masks_new_meta[p,3])) |
| c = as.integer(as.scalar(masks_new_meta[p,4])) |
| vec = masks_new[s, start:end] |
| previous_masked_b = matrix(vec, rows=r, cols=c, byrow=TRUE) |
| |
| # SANITY CHECK: dimensions with layers of previous layer match |
| if (l > 1 & ncol(previous_masked_b) != nrow(W)) { |
| if (verbose) print("W/prev layer mismatch in layer l=", l) |
| if (verbose) print("prev_b:", nrow(previous_masked_b), "x", ncol(previous_masked_b)) |
| if (verbose) print("W:", nrow(W), "x", ncol(W)) |
| stop("Invalid W shape wrt previous layer") |
| } |
| |
| if (nrow(previous_masked_b)==1) previous_masked_b = t(previous_masked_b) |
| if (ncol(masked_b) == 1) masked_b = t(masked_b) |
| |
| if(l==L) { # output layer |
| masked_W = previous_masked_b %*% matrix(1, 1, ncol(masked_W)) |
| } |
| else if (l<L & as.scalar(isFC[1, l+1]) == 0) { # next layer is not FC |
| masked_W = previous_masked_b %*% matrix(1, 1, ncol(masked_W)) |
| } else { |
| masked_W = previous_masked_b %*% masked_b |
| } |
| } |
| |
| # 4. forward these masks to all parameters in this layer |
| for (param in 1:paramsPerFCLayer) { |
| k = (param-1)*L + l |
| start = as.integer(as.scalar(masks_new_meta[k,1])) |
| end = as.integer(as.scalar(masks_new_meta[k,2])) |
| len = end - start + 1 |
| |
| if (param %% 2 == 0) { |
| flat = matrix(masked_b, rows=1, cols=len, byrow=TRUE) |
| } else { |
| flat = matrix(masked_W, rows=1, cols=len, byrow=TRUE) |
| } |
| masks_new[s, start:end] = flat |
| } |
| } |
| |
| # SANITY CHECK: accumulating all subnets bias's result in vector of all 1s (in FC hidden layers only) |
| if (l < L) { |
| if (as.scalar(isFC[1, l+1]) == 1) { |
| # bias parameter index for layer l is (L + l) |
| p = L + l |
| |
| start = as.integer(as.scalar(masks_new_meta[p,1])); |
| end = as.integer(as.scalar(masks_new_meta[p,2])); |
| r = as.integer(as.scalar(masks_new_meta[p,3])); |
| c = as.integer(as.scalar(masks_new_meta[p,4])); |
| |
| # sum across subnets for this bias slice |
| sumB_flat = colSums(masks_new[1:numSubnets, start:end]) |
| |
| # reshape back to bias shape (usually Hx1 or 1xH depending on how you store it) |
| sumB = matrix(sumB_flat, rows=r, cols=c, byrow=TRUE) |
| |
| if (min(sumB) != 1 | max(sumB) != 1) { |
| if (verbose) print("Subnet bias masks not a partition in layer l=" + l) |
| if (verbose) print("min(sumB)=" + min(sumB) + " max(sumB)=" + max(sumB)) |
| stop("Invalid subnet bias partition") |
| } |
| } |
| } |
| } |
| else { |
| # Non-FC layer: independent subnet training will not create disjoint partitions for this layer // shared across all subnets -> masks are all-ones |
| for (param in 1:paramsPerFCLayer) { |
| k = (param-1)*L + l |
| start = as.integer(as.scalar(masks_new_meta[k,1])) |
| end = as.integer(as.scalar(masks_new_meta[k,2])) |
| r = as.integer(as.scalar(masks_new_meta[k,3])) |
| c = as.integer(as.scalar(masks_new_meta[k,4])) |
| len = end - start + 1 |
| |
| if (param %% 2 == 0) { |
| shared_b = matrix(1, rows=r, cols=c) |
| flat = matrix(shared_b, rows=1, cols=len, byrow=TRUE) |
| } else { |
| shared_W = matrix(1, rows=r, cols=c) |
| flat = matrix(shared_W, rows=1, cols=len, byrow=TRUE) |
| } |
| masks_new[1:numSubnets, start:end] = matrix(1, numSubnets, 1) %*% flat |
| } |
| } |
| } |
| } |
