TensorFlow Variable Sequence Length

TensorFlow Variable Sequence Length

Overview

Variable Sequence Length (VSL) is a feature that allows computations on the CS system running in pipeline mode to process tensors which vary in shape from one element of a batch to the next.

In natural language processing applications, it is common for input data to consist of sequences of heterogeneous length. When this is the case, short samples get padded up to a user defined maximum sequence length so that they can be batched together. Naive treatment of padded data can result in wasting significant computation on padding tokens. VSL allows users to strip away this padding as samples enter the wafer and perform the model computation on non-padded variable length tensors. This leads to less wasted computation and faster training times. Typically models written for GPU include logic to ensure that padding tokens do not contribute to the final loss of the model. In this case, enabling VSL will have no affect on the model’s function other than increasing throughput.

Interface Details

VSL is enabled through a combination of using certain pad id values in input tensors and setting flags in the Cerebras configuration according to the following rules.

  • You must set config.matching.kernel.use_legacy_vsl = True in the Cerebras configuration.

  • If an input tensor is first used in an embedding layer, then that input needs to be padded with a value that isn’t used for any purpose other than padding, and that padding value must be supplied to the Cerebras embedding layer implementation. Supplying a pad value of -1 to the attention layer signifies that VSL should not be used. When the embedding kernel sees a pad value other than -1, it strips away all tokens at the end of the input sequence it receives that start with the pad value. For example,if you want to enable VSL in a BERT model, then you could pad the input_ids tensor with -100 and specify to the embedding layer that the pad_id is -100. You would need to do similar for the other inputs to the model that are fed directly to embedding layers.

  • For a labels tensor fed directly into an MLM loss computation, this tensor must be padded with -1.

As long as the pad_id is correctly specified in any embedding layers, there is no change necessary to model code to enable VSL, only the above changes to data and the Cerebras config.

Limitations

VSL is limited in its generality and is currently in the process of being replaced by VTS, a new, more general version of the same functionality with a cleaner interface. See PyTorch Variable Tensor Shape for more details.

As mentioned above, VSL is limited in its generality. Accordingly, if a model does not compile with VSL turned on, we suggest attempting a compile without VSL. This feature is also prone to user bugs involving improperly set pad ids or not fully enabling VSL for all inputs of the model. If the pad_id specified is used in some part of the input other than padding or if the lengths of input tensors or label tensors after removing padding don’t align, the model will compile without error but will stall during runtime.