System Info
transformers version: 4.37.0.dev0- Platform: macOS-13.5-arm64-arm-64bit
- Python version: 3.10.13
- Huggingface_hub version: 0.20.1
- Safetensors version: 0.4.1
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.1.2 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?:
- Using distributed or parallel set-up in script?:
Who can help?
@ArthurZucker
Information
Tasks
Reproduction
Two issues were found:
mixtral's implementation of auxiliary loss is not correct.
I think load_balancing_loss_func in modeling_mixtral computes auxiliary loss incorrectly
|
def load_balancing_loss_func(gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2) -> float: |
|
r""" |
|
Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch. |
|
|
|
See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss |
|
function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between |
|
experts is too unbalanced. |
|
|
|
Args: |
|
gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]): |
|
Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of |
|
shape [batch_size X sequence_length, num_experts]. |
|
num_experts (`int`, *optional*): |
|
Number of experts |
|
|
|
Returns: |
|
The auxiliary loss. |
|
""" |
|
if gate_logits is None or not isinstance(gate_logits, tuple): |
|
return 0 |
|
|
|
if isinstance(gate_logits, tuple): |
|
compute_device = gate_logits[0].device |
|
concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0) |
|
|
|
routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1) |
|
|
|
_, selected_experts = torch.topk(routing_weights, top_k, dim=-1) |
|
|
|
# treat `top_k` as tokens (shape is `top_k X [batch_size X sequence_length]`) |
|
selected_experts = selected_experts.reshape(-1) |
|
|
|
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) |
|
expert_mask = torch.max(expert_mask, dim=-2).values |
|
|
|
# Compute the percentage of tokens routed to each experts |
|
tokens_per_expert = torch.mean(expert_mask.float(), dim=0) |
|
|
|
# Compute the average probability of routing to these experts |
|
router_prob_per_expert = torch.mean(routing_weights, dim=0) |
|
|
|
overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(-1)) |
|
return overall_loss * num_experts |
Auxiliary loss is implemented as multiply fraction of tokens dispatched to expert by fraction of the router probability allocated for expert. The fraction of tokens dispatched to expert is calculated as the number of tokens routed to expert divided by the total number of tokens. The actual implementation is as follows:
|
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts) |
|
expert_mask = torch.max(expert_mask, dim=-2).values |
|
|
|
# Compute the percentage of tokens routed to each experts |
|
tokens_per_expert = torch.mean(expert_mask.float(), dim=0) |
As we know, the shape of
selected_experts is
top_k X [batch_size X sequence_length]，so the shape of
expert_mask is
[top_k X batch_size X sequence_length, num_experts]. When we excute
expert_mask = torch.max(expert_mask, dim=-2).values , the shape of the
expert_mask becomes
[num_experts], which means that whenever a token is routed to an expert, that expert has a value of 1. After the operation of
torch.mean(expert_mask.float(), dim=0),
tokens_per_expert becomes a scaler, which is clearly incorrect, since tokens_per_expert should have a shape of
[num_experts].
Example
Example Inputs：
T = 3 # number of tokens [B X S]
num_experts = 8
top_k = 2 # top_2
gate_logits = torch.randn(T,num_experts)
routing_weights = torch.nn.functional.softmax(gate_logits, dim=-1)
Each row of routing_weights represents the probability that a token will be routed to an expert
tensor([[0.2551, 0.2519, 0.0357, 0.0830, 0.0897, 0.0981, 0.1565, 0.0299],
[0.0728, 0.0593, 0.0948, 0.1708, 0.0098, 0.0848, 0.3884, 0.1192],
[0.0292, 0.0387, 0.0696, 0.1331, 0.6699, 0.0049, 0.0442, 0.0104]])next select experts
_, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
# treat `top_k` as tokens (shape is `top_k X [batch_size X sequence_length]`)
selected_experts = selected_experts.reshape(-1)
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
we get the following result (shape [top_k X batch_size X sequence_length, num_experts]):
tensor([[1, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0, 0]])The final results are as follows
expert_mask = torch.max(expert_mask, dim=-2).values
# tensor([1, 1, 0, 1, 1, 0, 1, 0])
tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
# tensor(0.6250)
router_prob_per_expert = torch.mean(routing_weights, dim=0)
# tensor([0.0746, 0.1031, 0.0448, 0.0804, 0.1823, 0.1216, 0.3112, 0.0820])
overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(-1))
# tensor(0.6250)
Because the sum of the router_prob_per_expert is 1, the final loss value is actually the value of tokens_per_expert. As the total number of tokens increases, the value of tokens_per_expert will be 1 (each expert has tokens routed to).
Solution
The tokens_per_expert calculation should divide the tokens that are routed per expert by the total number of tokens. Specifically, we can sum the columns of expert_mask and divide by the total number of tokens. The following is an implementation
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
expert_mask = expert_mask.reshape(-1, top_k, num_experts)
expert_mask = torch.max(expert_mask, dim=-2).values
# Compute the percentage of tokens routed to each experts
tokens_per_expert = torch.mean(expert_mask.float(), dim=0) / top_k
# Compute the average probability of routing to these experts
router_prob_per_expert = torch.mean(routing_weights, dim=0)
overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert)
return overall_loss * num_experts
Example
expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
expert_mask = expert_mask.reshape(T,top_k,-1)
'''
tensor([[[1, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0, 0]],
[[0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 1, 0, 0, 0, 0]],
[[0, 0, 0, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 0, 0]]])
'''
expert_mask = torch.max(expert_mask, dim=-2).values
'''
tensor([[1, 1, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 1, 0, 0, 1, 0],
[0, 0, 0, 1, 1, 0, 0, 0]])
'''
tokens_per_expert = torch.mean(expert_mask.float(), dim=0) / top_k
# tensor([0.1667, 0.1667, 0.0000, 0.3333, 0.1667, 0.0000, 0.1667, 0.0000])Note
On the other hand, in switch transformer ((https://arxiv.org/abs/2101.03961), auxiliary loss should converge to 1 when the load is balanced. However, the top 1 strategy is used in the paper, so the maximum value is taken when calculating tokens_per_expert. In the top_k strategy, this corresponds to top_k*T tokens being routed to the experts, so tokens_per_expert should be divided by top_k. Otherwise the final converged value should be top_k.
By the way, the unit test should determine if the loss is close to 1 instead of 8.
|
torch.testing.assert_close(result.aux_loss.cpu(), torch.tensor(8, dtype=torch.float32)) |
Should the output of each layer of the gated network be concatenated into a tensor?
https://github.com/huggingface/transformers/blob/3cefac1d974db5e2825a0cb2b842883a628be7a0/src/transformers/models/mixtral/modeling_mixtral.py#L98C5-L101C1
Before calculating the auxiliary loss, the routing outputs of the different transformer layers of the expert layer are concatenated into a tensor. This implies that the routing outputs of different layers are mapped to the same expert, and in fact the routing outputs of each layer should be mapped to its own layer of experts. So should the auxiliary loss be calculated for each layer independently, rather than concatenated into a tensor？
Expected behavior
I expect to examine the problem and review my solution for the first issue and have a discussion about the second issue, as I'm not sure if it makes more sense to calculate loss separately.
System Info
transformersversion: 4.37.0.dev0Who can help?
@ArthurZucker
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
Two issues were found:
mixtral's implementation of auxiliary loss is not correct.
I think
load_balancing_loss_funcinmodeling_mixtralcomputes auxiliary loss incorrectlytransformers/src/transformers/models/mixtral/modeling_mixtral.py
Lines 77 to 119 in 3cefac1
Auxiliary loss is implemented as multiply fraction of tokens dispatched to expert by fraction of the router probability allocated for expert. The fraction of tokens dispatched to expert is calculated as the number of tokens routed to expert divided by the total number of tokens. The actual implementation is as follows:
transformers/src/transformers/models/mixtral/modeling_mixtral.py
Lines 109 to 113 in 3cefac1
As we know, the shape of
selected_expertsistop_k X [batch_size X sequence_length]，so the shape ofexpert_maskis[top_k X batch_size X sequence_length, num_experts]. When we excuteexpert_mask = torch.max(expert_mask, dim=-2).values, the shape of theexpert_maskbecomes[num_experts], which means that whenever a token is routed to an expert, that expert has a value of 1. After the operation oftorch.mean(expert_mask.float(), dim=0),tokens_per_expertbecomes a scaler, which is clearly incorrect, since tokens_per_expert should have a shape of[num_experts].Example
Example Inputs：
Each row of
routing_weightsrepresents the probability that a token will be routed to an expertnext select experts
we get the following result (shape
[top_k X batch_size X sequence_length, num_experts]):The final results are as follows
Because the sum of the
router_prob_per_expertis 1, the final loss value is actually the value oftokens_per_expert. As the total number of tokens increases, the value oftokens_per_expertwill be 1 (each expert has tokens routed to).Solution
The
tokens_per_expertcalculation should divide the tokens that are routed per expert by the total number of tokens. Specifically, we can sum the columns ofexpert_maskand divide by the total number of tokens. The following is an implementationExample
Note
On the other hand, in switch transformer ((https://arxiv.org/abs/2101.03961), auxiliary loss should converge to 1 when the load is balanced. However, the top 1 strategy is used in the paper, so the maximum value is taken when calculating tokens_per_expert. In the top_k strategy, this corresponds to
top_k*Ttokens being routed to the experts, so tokens_per_expert should be divided bytop_k. Otherwise the final converged value should betop_k.By the way, the unit test should determine if the loss is close to 1 instead of 8.
transformers/tests/models/mixtral/test_modeling_mixtral.py
Line 477 in 3cefac1
Should the output of each layer of the gated network be concatenated into a tensor?
https://github.com/huggingface/transformers/blob/3cefac1d974db5e2825a0cb2b842883a628be7a0/src/transformers/models/mixtral/modeling_mixtral.py#L98C5-L101C1
Before calculating the auxiliary loss, the routing outputs of the different transformer layers of the expert layer are concatenated into a tensor. This implies that the routing outputs of different layers are mapped to the same expert, and in fact the routing outputs of each layer should be mapped to its own layer of experts. So should the auxiliary loss be calculated for each layer independently, rather than concatenated into a tensor？
Expected behavior
I expect to examine the problem and review my solution for the first issue and have a discussion about the second issue, as I'm not sure if it makes more sense to calculate loss separately.