Debugging sft_student mode

[karthikviswanathn]

This issue is based on the notebook here. I was comparing the next token prediction (NTP) on the student mode using set_transformer_early_exit_mode(model, 'sft_student') and manual prediction using exiting and freezing usig KV caching.

1. When there were no early exits, the NTP was identical. This is expected and I was hoping it would extend to any general exiting.
2. When I force an early exit on a single token, the NTP using the student mode is gibberish for later tokens whereas the manual prediction gives sane predictions only slightly different from the base model.

Maybe there is some unexpected behaviour is going on in patched_layer_forward() and patched_attention_forward()?

[karthikviswanathn]

The student output looks sane now as we can see in the notebook here. However the following part gives incorrect hidden states.

student_outputs = model(repeated_sft_teacher_generated_tokens, prescribed_exit_layer_idxs = chosen_exit_layers_tensor, output_hidden_states = True)    
student_hidden_states = student_outputs[0].hidden_states

The cause of this issue seems to be due to operations like hidden_states[unfrozen_elements]. This kind of indexing disrupts the underlying tensor references, which prevents the transformer's hidden states from being updated as expected. For example,

hidden_states[unfrozen_elements] = self.post_attention_layernorm(hidden_states[unfrozen_elements])

can be replaced by

hidden_states = torch.where(unfrozen_elements.unsqueeze(-1), self.post_attention_layernorm(hidden_states), _original_hidden_states)

The exact details have to be figured out though. This can be useful when we are testing for the hidden states used to make the decision weights.