from easy_context import prepare_seq_parallel_inputs, apply_seq_parallel_monkey_patch, prepare_dataloaderfrom transformers import LlamaForCausalLM# Swap attention implementation from flash attn to either dist_ring_attn or zigzag_ring_attnapply_seq_parallel_monkey_patch("dist_flash_attn", "llama")# Make sure you toggle on flash_attention_2model = LlamaForCausalLM.from_pretrained(model_name, _attn_implementation="flash_attention_2")accelerator = ...train_dataloader = ...prepare_dataloader("dist_flash_attn", train_dataloader, accelerator)
# In your training loop...for step, batch in enumerate(train_dataloader):# Shard the sequencesprepared = prepare_seq_parallel_inputs("dist_flash_attn", batch["input_ids"], batch["position_ids"], batch["target_ids"], accelerator.process_index, accelerator.num_processes, accelerator.device)local_input_ids = prepared["local_input_ids"]local_position_ids = prepared["local_position_ids"]local_target_ids = prepared["local_target_ids"]# Then do model forward as usuallogits = model(local_input_ids,position_ids=local_position_ids,).logits