Feature request
Currently the layer_idx is recorded in the attention module of each LlamaDecoderLayer. This has the unfortunate side effect that the layers cannot easily be moved around or reused within the layer list. It seems simple enough to pass in the layer index as part of loop over layers in the forward pass. That way the layers once again will be decouple from their position information.
Backward compatibility could be preserved by still accepting the argument in the constructor but defaulting it to None and then just ignoring it in favor of the passed forward argument.
Motivation
The motivation is to allow for simple layer stacking (like we have been seeing with pass through merged models) at inference time without actually expanding the memory usage of the model.
Your contribution
I am happy to send a PR. Seems simple enough.
Feature request
Currently the layer_idx is recorded in the attention module of each
LlamaDecoderLayer. This has the unfortunate side effect that the layers cannot easily be moved around or reused within the layer list. It seems simple enough to pass in the layer index as part of loop over layers in the forward pass. That way the layers once again will be decouple from their position information.Backward compatibility could be preserved by still accepting the argument in the constructor but defaulting it to None and then just ignoring it in favor of the passed forward argument.
Motivation
The motivation is to allow for simple layer stacking (like we have been seeing with pass through merged models) at inference time without actually expanding the memory usage of the model.
Your contribution
I am happy to send a PR. Seems simple enough.