sequence parallel fix bug (#47)

* feat(tests): add sequence parallel single attention test

Add a new test file `test_sequence_parallel_single_attention.py` to verify the correctness of the sequence parallel attention implementation. The test includes a distributed setup using torch.distributed and compares outputs between sequence parallel and local attention modes. Also adds an empty `__init__.py` to the transformers test directory for proper module imports.

* wip

* feat(tests): enhance sequence parallel attention test determinism

- Add `_enable_strict_determinism` helper to disable TF32 and enable deterministic algorithms
- Add `_to_local` helper to unwrap DTensors for gradient comparison
- Update test to use full world size for sequence parallel group and increase head count
- Switch to float32 dtype for stricter numerical alignment
- Improve gradient comparison by cloning and unwrapping tensors

* remove __init__

* feat(sequence_parallel): refactor config handling and remove padding-free logic

- Replace HfConfigFactory utility with direct get_config_attr function
- Move get_llm_model to shared transformers utilities
- Remove padding_free parameter and related conditional logic
- Simplify attention mask construction for padded tokens
- Update SequenceParallelConfig to drop padding_free field

* feat(sequence_parallel): enforce flash_attention_2 for packed batches

- Add detection of packed batches via `_is_packed_position_ids` heuristic
- Raise RuntimeError when SDPA backend is used with packed batches, as SDPA lacks native packed/varlen support
- Build 2D attention_mask for padded sequences to ensure correct FlashAttention2 unpad behavior
- Avoid unnecessary 4D causal mask generation for packed/padding-free batches

* feat(sft): add single controller SP packing example for Qwen2.5-7B

Introduce a new cookbook script demonstrating supervised fine-tuning with a single controller using sequence parallelism (SP) and FSDP across 4 GPUs. The example includes:
- Device mesh configuration with dp=2 and fsdp=2 dimensions
- PackingDataset setup with self-cognition data and left truncation
- Training loop with LoRA adapter, AdamW optimizer, and periodic evaluation
- Checkpoint saving based on loss improvement
- Validation of FSDP + SP input slicing across multiple GPUs

* fix loss computation bug

* feat(cookbook): add single controller SP example and reorganize transformers cookbook

- Add new single_controller_sp.py example demonstrating FSDP + SP validation over 4 GPUs
- Move legacy single_controller_sp.py to transformers/sp_fsdp_dense.py
- Add shell script sp_fsdp_dense.sh for running the example
- Update imports and structure to use twinkle framework components

* refactor(tests): move sequence parallel attention test to dedicated directory

Relocate test_sequence_parallel_single_attention.py from tests/transformers/ to tests/sequence_parallel/ to better organize test files by feature area. This improves maintainability and aligns with the project's test structure conventions.

* feat: add sequence parallelism instructions and clean up imports

- Add bash script header and comments to `sp_fsdp_dense.sh` explaining how to enable sequence parallelism with ulysses_size
- Remove duplicate `import os` statement in transformers.py for cleaner code
- Fix minor formatting by removing extra blank line in transformers_utils.py

* refactor

* feat: update training script with local mode and evaluation

- Switch from `ray` to `local` mode for twinkle initialization
- Add evaluation function with separate dataset slice
- Increase dataset size from 100 to 500 samples
- Add cosine warmup learning rate scheduler
- Remove unused torch import and remote_group parameters
- Adjust batch size from 4 to 8 and logging frequency to every 20 steps
- Improve logging with train configs and total steps information

* feat(transformers): remove unused imports in sequence_parallel module

Removed unnecessary imports (`math`, `os`, `SimpleNamespace`) from the sequence_parallel strategy file to clean up the codebase and improve maintainability.

Author: meichangsu1

cookbook/legacy/single_controller_sp.py cookbook/transformers/sp_fsdp_dense.pycookbook/legacy/single_controller_sp.py renamed to cookbook/transformers/sp_fsdp_dense.py

@@ -1,6 +1,5 @@
 from functools import partial
 import numpy as np
-import torch
 from peft import LoraConfig
 
 import twinkle
@@ -12,6 +11,7 @@
 
 logger = get_logger()
 MODEL_ID = 'ms://Qwen/Qwen2.5-7B-Instruct'
+DATASETS='ms://swift/self-cognition'
 
 device_group = [
     DeviceGroup(
@@ -30,79 +30,70 @@
 )
 
 twinkle.initialize(
-    mode="ray",
+    mode="local",
     nproc_per_node=4,
-    groups=device_group,
     global_device_mesh=device_mesh,
     lazy_collect=False,
 )
 
+def eval(model):
+    dataloader = DataLoader(
+        dataset=partial(create_dataset, data_slice=range(100)),
+        batch_size=4,
+        device_mesh=device_mesh,
+    )
+    for _, batch in enumerate(dataloader):
+        model.forward_only(inputs=batch, adapter_name="default")
+        model.calculate_loss(adapter_name="default")
+    return model.calculate_metric(is_training=False, adapter_name="default")
+
 
 def create_dataset(data_slice=None):
     dataset = Dataset(
-        dataset_meta=DatasetMeta("ms://swift/self-cognition", data_slice=data_slice)
+        dataset_meta=DatasetMeta(DATASETS, data_slice=range(500))
     )
     dataset.set_template(
         "Template",
-        model_id=MODEL_ID,
-        truncation_strategy="left",
-        max_length=64,
+        model_id=MODEL_ID
     )
     dataset.map(SelfCognitionProcessor("twinkle模型", "twinkle团队"))
     dataset.encode(batched=True)
     return dataset
-
-
-def eval(model: TransformersModel):
-    dataloader = DataLoader(
-        dataset=partial(create_dataset, data_slice=range(20)),
-        batch_size=4,
-        drop_last=True,
-        device_mesh=device_mesh,
-        remote_group="default",
-    )
-    for step, batch in enumerate(dataloader):
-        model.forward_only(inputs=batch, adapter_name="default")
-        model.calculate_loss(adapter_name="default")
-    metrics = model.calculate_metric(is_training=False, adapter_name="default")
-    return metrics()
-
-
 def train():
     dataloader = DataLoader(
         dataset=partial(create_dataset, data_slice=None),
-        batch_size=4,
+        batch_size=8,
         device_mesh=device_mesh,
-        remote_group="default",
     )
 
     model = TransformersModel(
         model_id=MODEL_ID,
         device_mesh=device_mesh,
         strategy="native_fsdp",
-        remote_group="default",
     )
 
     lora_config = LoraConfig(target_modules="all-linear")
     model.add_adapter_to_model("default", lora_config, gradient_accumulation_steps=1)
     model.set_optimizer("AdamW", lr=1e-4, adapter_name="default")
+    model.set_lr_scheduler(
+        scheduler_cls="CosineWarmupScheduler",
+        num_warmup_steps=5,
+        num_training_steps=len(dataloader),
+        adapter_name="default",
+    )
+
+    logger.info(model.get_train_configs(adapter_name="default"))
+    logger.info(f"Total steps: {len(dataloader)}")
+
 
-    loss_metric = 99.0
     for step, batch in enumerate(dataloader):
-        if isinstance(batch, list) and len(batch) == 0:
-            continue
-        output = model.forward_backward(inputs=batch, adapter_name="default")
-        loss_value = output() if callable(output) else output
-        logger.info(f"step {step}, loss: {loss_value}")
+        model.forward_backward(inputs=batch, adapter_name="default")
         model.clip_grad_and_step(adapter_name="default")
-        if step % 50 == 0 and step > 0:
-            metrics = eval(model)
-            logger.info(f"Current is step {step} of {len(dataloader)}, metric: {metrics}")
-            metrics["step"] = step
-            if loss_metric > metrics["loss"]:
-                model.save(f"checkpoint-{step}")
-                loss_metric = metrics["loss"]
+        if step % 20 == 0:
+            metric = model.calculate_metric(is_training=True, adapter_name="default")
+            logger.info(f"Current is step {step} of {len(dataloader)}, metric: {metric}")
+    model.save("last-checkpoint", interval=1)
 
 
 if __name__ == "__main__":
-    train()
+    train()

cookbook/transformers/sp_fsdp_dense.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+# To enabele sequence parallelism, please set ulysses_size > 1
+# device_mesh = DeviceMesh(
+#     device_type="cuda",
+#     mesh=np.arange(4).reshape(2, 2),
+#     mesh_dim_names=("dp", "fsdp"),
+#     ulysses_size=2,
+# )
+#
+CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 sp_fsdp_dense.py