What:

Ensures tensors are contiguous in memory with matching strides during the backward pass.

Fixes #40

Why:

Multiple users/customers have been facing issues while training with the axolotl package see #40

maybe_contiguous fails to adequately check if tensors are contiguous in certain packing scenarios

Changes:

• flash_attn_interface.py
  • _flash_attn_backward: stride check, contiguous check
  • _flash_attn_varlen_backward: contiguous check

Testing:

• Manual testing: Successfully fine tuned Phi-2, StableLM, and TinyLlama in axolotl
• Automated testing: Benchmarking scripts ran as expected