layer number bug fixes

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.

[ksivaman]

Fixes#294, UnfusedDotProductAttention for fp16, and KV inference cache.

[timmoon10]

LGTM, I don't see any behavior change as long as a valid layer_number is provided.

That said, can you clarify the behavior of layer_number? It's still a bit opaque to me. From what I can see:

• It is used as a key in the KV inference cache:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 1151 to 1171 in 706d7da

  	is_first_step=False
  	ifinference_paramsandself.layer_numberisnotNone:
  	ifself.layer_numbernotininference_params.key_value_memory_dict:
  	inf_max_seq_len=inference_params.max_sequence_len
  	inf_max_batch_size=inference_params.max_batch_size
  	inference_key_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_value_memory=self._allocate_memory(
  	inf_max_seq_len, inf_max_batch_size, hidden_states.dtype
  	)
  	inference_params.key_value_memory_dict[self.layer_number] = (
  	inference_key_memory,
  	inference_value_memory,
  	)
  	is_first_step=True
  	else:
  	(
  	inference_key_memory,
  	inference_value_memory,
  	) =inference_params.key_value_memory_dict[self.layer_number]

  
  For this role, layer_number can be any hashable type as long as it is unique to each attention module.
• It is a scaling factor to prevent FP16 overflows when computing attention scores:
  

  TransformerEngine/transformer_engine/pytorch/attention.py

  Lines 196 to 198 in 706d7da

  	scale=self.norm_factor
  	ifapply_qk_layer_scaling:
  	scale*=self.layer_number

  
  I don't understand this though. Each value is the dot product of two (hidden_dim/num_heads) vectors, so I'd imagine the values would be independent of the number of layers, sequence length, batch size, etc.

[ksivaman]

/te-ci

[ksivaman]

@timmoon10 in the second scenario, layer_number is used as a scaling technique for better convergence with fp16 training.

[timmoon10]

Can you give some intuition or documentation? All layers should have QKV values around the same range due to layernorm, so the QK GEMM should also have a consistent range and not require scaling. It might make sense in bprop if there is some growth or decay in the gradients, but why do we scale by layer_number instead of 1/layer_number or exp(+/- layer_number) (the last seems most natural if we have a exploding/vanishing gradients problem)? I'm increasingly suspicious of combining the KV cache index with the scaling factor since there's no reason they have to be related.

Edit: Based on discussions, it seems this is a trick from Megatron-LM. The scaling factors are a hyperparameter knob to explore in the future.