# Source code for fairseq.modules.base_layer

```# Copyright (c) Facebook, Inc. and its affiliates.
#
# LICENSE file in the root directory of this source tree.

import torch.nn as nn
import torch
import sys
from fairseq import utils
from fairseq.distributed import utils as distributed_utils
from fairseq.modules.layer_norm import LayerNorm

[docs]class BaseLayer(nn.Module):
def __init__(self, args):
super().__init__()
self.num_workers = distributed_utils.get_data_parallel_world_size()
expert_centroids = torch.empty(self.num_workers, args.decoder_embed_dim)
torch.nn.init.orthogonal_(expert_centroids, gain=0.1)
self.register_parameter(
"expert_centroids", torch.nn.Parameter(expert_centroids)
)
self.expert_network = nn.Sequential(
*([BaseSublayer(args) for _ in range(args.base_sublayers)])
)
self.expert_id = distributed_utils.get_data_parallel_rank()
self.shuffle = args.base_shuffle

# Add a special attribute to the expert parameters, so we know not to sync their gradients
for param in self.expert_network.parameters():
param.expert = True

[docs]    def forward(self, input_features, *args, **kwargs):
features = input_features.reshape(-1, input_features.size(-1))

if self.shuffle and is_training:
# Send each token to a random worker, to break correlations within the batch
shuffle_sort = torch.randperm(features.size(0), device=features.device)
features = All2All.apply(features[shuffle_sort])

# Compute similarity of each token to each expert, for routing
token_expert_affinities = features.matmul(
self.expert_centroids.transpose(0, 1)
)

# Compute which token goes to which expert
sort_by_expert, input_splits, output_splits = (
self.balanced_assignment(token_expert_affinities)
if is_training
else self.greedy_assignment(token_expert_affinities)
)
# Swap these tokens for the right ones for our expert
routed_features = All2All.apply(
features[sort_by_expert], output_splits, input_splits
)

if routed_features.size(0) > 0:
# Mix in the expert network based on how appropriate it is for these tokens
alpha = torch.sigmoid(
routed_features.mv(self.expert_centroids[self.expert_id])
).unsqueeze(1)
routed_features = (
alpha * self.expert_network(routed_features)
+ (1 - alpha) * routed_features
)
result = All2All.apply(routed_features, input_splits, output_splits)[
self.inverse_sort(sort_by_expert)
]

if self.shuffle and is_training:
# Undo shuffling
result = All2All.apply(result)[self.inverse_sort(shuffle_sort)]

# Return additional Nones for compatibility with TransformerDecoderLayer
return result.view(input_features.size()), None, None

[docs]    def inverse_sort(self, order):
# Creates an index that undoes a sort: xs==xs[order][inverse_sort(order)]
0, order, torch.arange(0, order.size(0), device=order.device)
)

[docs]    def balanced_assignment(self, scores):
ok = scores.isfinite()
if not ok.all():
# NaNs here can break the assignment algorithm
scores[~ok] = scores[ok].min()
return self.cpp.balanced_assignment(scores), None, None

# Assigns each token to the top k experts
[docs]    def greedy_assignment(self, scores, k=1):
token_to_workers = torch.topk(scores, dim=1, k=k, largest=True).indices.view(-1)
token_to_workers, sort_ordering = torch.sort(token_to_workers)
worker2token = sort_ordering // k

# Find how many tokens we're sending to each other worker (being careful for sending 0 tokens to some workers)
output_splits = torch.zeros(
(self.num_workers,), dtype=torch.long, device=scores.device
)
workers, counts = torch.unique_consecutive(token_to_workers, return_counts=True)
output_splits[workers] = counts
# Tell other workers how many tokens to expect from us
input_splits = All2All.apply(output_splits)
return worker2token, input_splits.tolist(), output_splits.tolist()

try:
from fairseq import libbase

return libbase

except ImportError as e:
sys.stderr.write(
"ERROR: missing libbase. run `python setup.py build_ext --inplace`\n"
)
raise e

class BaseSublayer(nn.Module):
def __init__(self, args):
super().__init__()
self.activation_fn = utils.get_activation_fn(
activation=getattr(args, "activation_fn", "relu") or "relu"
)
self.norm = LayerNorm(args.decoder_embed_dim, export=False)
self.ff1 = torch.nn.Linear(args.decoder_embed_dim, args.decoder_ffn_embed_dim)
self.ff2 = torch.nn.Linear(args.decoder_ffn_embed_dim, args.decoder_embed_dim)
self.ff2.weight.data.zero_()

def forward(self, xs):
return xs + self.ff2(self.activation_fn(self.ff1(self.norm(xs))))

# Wraps torch.distributed.all_to_all_single as a function that supports autograd
@staticmethod
def forward(ctx, xs, input_splits=None, output_splits=None):
ctx.input_splits = input_splits
ctx.output_splits = output_splits

ys = (
torch.empty_like(xs)
if output_splits is None
else xs.new_empty(size=[sum(output_splits)] + list(xs.size()[1:]))
)
torch.distributed.all_to_all_single(
ys, xs, output_split_sizes=output_splits, input_split_sizes=input_splits
)
return ys

@staticmethod
result = (
{"ft_lang_label":"__label__en","ft_lang_prob":0.53051096,"math_prob":0.9114187,"size":5492,"snap":"2023-14-2023-23","text_gpt3_token_len":1298,"char_repetition_ratio":0.14996356,"word_repetition_ratio":0.0036968577,"special_character_ratio":0.2414421,"punctuation_ratio":0.24363233,"nsfw_num_words":0,"has_unicode_error":false,"math_prob_llama3":0.95425725,"pos_list":[0],"im_url_duplicate_count":[null],"WARC_HEADER":"{\"WARC-Type\":\"response\",\"WARC-Date\":\"2023-06-08T00:57:44Z\",\"WARC-Record-ID\":\"<urn:uuid:e8449884-736a-4362-b4f6-e72f3bee0027>\",\"Content-Length\":\"39266\",\"Content-Type\":\"application/http; msgtype=response\",\"WARC-Warcinfo-ID\":\"<urn:uuid:34d461d1-b0c8-4026-b5ad-a2851bad4957>\",\"WARC-Concurrent-To\":\"<urn:uuid:c702b1e0-0a7f-4bda-8ec3-a50a85e6a634>\",\"WARC-IP-Address\":\"104.17.32.82\",\"WARC-Target-URI\":\"https://fairseq.readthedocs.io/en/latest/_modules/fairseq/modules/base_layer.html\",\"WARC-Payload-Digest\":\"sha1:SWSABFXEOQWMP3AO7NKEM4J55NNDPGC4\",\"WARC-Block-Digest\":\"sha1:ERYVJAOXKV3SD2H524SZAAPZRYZDNWX2\",\"WARC-Identified-Payload-Type\":\"text/html\",\"warc_filename\":\"/cc_download/warc_2023/CC-MAIN-2023-23/CC-MAIN-2023-23_segments_1685224654031.92_warc_CC-MAIN-20230608003500-20230608033500-00131.warc.gz\"}"}