[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

examples/pytorch/transformer/context_parallel_runner_bshd.py

@@ -14,7 +14,9 @@
 
 from utils import get_dummy_data_bshd, collect_gradients, DistributedConfig
 from model import SimpleConfig, SimpleBSHDModel
-from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import get_batch_on_this_cp_rank
+from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
+    get_batch_on_this_cp_rank,
+)
 
 
 random.seed(42)
@@ -23,12 +25,12 @@
 if torch.cuda.is_available():
     torch.cuda.manual_seed_all(42)
 
-logging.basicConfig(level=logging.INFO, format='%(message)s')
+logging.basicConfig(level=logging.INFO, format="%(message)s")
 logger = logging.getLogger(__name__)
 
 data = get_dummy_data_bshd()
 
-DISTRIBUTED_MODE = 'RANK' in os.environ and 'WORLD_SIZE' in os.environ
+DISTRIBUTED_MODE = "RANK" in os.environ and "WORLD_SIZE" in os.environ
 
 # STEP 1: RUN CP=1 (BASELINE) - NO DISTRIBUTED TRAINING
 
@@ -55,7 +57,10 @@
     torch.cuda.manual_seed(42)
 
 batch = get_dummy_data_bshd()
-batch = {k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+batch = {
+    k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v
+    for k, v in batch.items()
+}
 
 with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
     output_cp1 = model_cp1(batch)
@@ -64,34 +69,36 @@
 labels_flat = batch["labels"].view(-1)
 valid_mask = labels_flat != -100
 if valid_mask.any():
-    loss_cp1 = F.cross_entropy(logits_flat_cp1[valid_mask], labels_flat[valid_mask], reduction="mean")
+    loss_cp1 = F.cross_entropy(
+        logits_flat_cp1[valid_mask], labels_flat[valid_mask], reduction="mean"
+    )
 else:
     loss_cp1 = logits_flat_cp1.sum() * 0.0
 
 # Compute gradients for CP=1
 loss_cp1.backward()
 
 target_layers = [
-    'embedding',           # Embedding layers
-    'transformer_layers.0',      # First transformer layer
-    'transformer_layers.1',      # Second transformer layer  
-    'linear'              # Language model head
+    "embedding",  # Embedding layers
+    "transformer_layers.0",  # First transformer layer
+    "transformer_layers.1",  # Second transformer layer
+    "linear",  # Language model head
 ]
 
 grads_cp1 = collect_gradients(model_cp1, layer_patterns=target_layers, max_params=15)
 
 initial_state_dict = {k: v.cpu().clone() for k, v in model_cp1.state_dict().items()}
-torch.save(initial_state_dict, '/tmp/bshd_initial_model_state.pt')
+torch.save(initial_state_dict, "/tmp/bshd_initial_model_state.pt")
 
 cp1_results = {
-    'logits': output_cp1.clone().detach().cpu(),
-    'loss': loss_cp1.clone().detach().cpu(),
-    'grad_norms': {name: grad.norm().item() for name, grad in grads_cp1.items()},
-    'grads': grads_cp1,
+    "logits": output_cp1.clone().detach().cpu(),
+    "loss": loss_cp1.clone().detach().cpu(),
+    "grad_norms": {name: grad.norm().item() for name, grad in grads_cp1.items()},
+    "grads": grads_cp1,
 }
 
-torch.save(cp1_results, '/tmp/bshd_cp1_results.pt')
-torch.save(data, '/tmp/bshd_data.pt')
+torch.save(cp1_results, "/tmp/bshd_cp1_results.pt")
+torch.save(data, "/tmp/bshd_data.pt")
 if not DISTRIBUTED_MODE:
     logger.info(f"CP=1 complete: loss={cp1_results['loss'].item():.6f}")
 
@@ -100,26 +107,27 @@
 # Skip CP=2 if not in distributed mode
 if not DISTRIBUTED_MODE:
     import sys
+
     sys.exit(0)
 
 # Run CP=2 in distributed mode
 if DISTRIBUTED_MODE:
-    if int(os.environ.get('RANK', 0)) == 0:
+    if int(os.environ.get("RANK", 0)) == 0:
         logger.info("Running CP=2 distributed...")
 
     cp_size = 2
     dist.init_process_group(backend="nccl")
     dist_config = DistributedConfig()
     torch.cuda.set_device(dist_config.local_rank)
     device_mesh = init_device_mesh(
-            "cuda",
-            mesh_shape=(1, cp_size, 1),
-            mesh_dim_names=("fsdp", "cp", "tp"),
-        )
+        "cuda",
+        mesh_shape=(1, cp_size, 1),
+        mesh_dim_names=("fsdp", "cp", "tp"),
+    )
     device = torch.device(f"cuda:{dist_config.local_rank}")
 
     model = SimpleBSHDModel(config)
-    initial_state_dict = torch.load('/tmp/bshd_initial_model_state.pt', map_location='cpu')
+    initial_state_dict = torch.load("/tmp/bshd_initial_model_state.pt", map_location="cpu")
     model.load_state_dict(initial_state_dict)
 
     model = model.to(device)
@@ -138,7 +146,7 @@
             transformer_layer.set_context_parallel_group(
                 cp_group,
                 torch.distributed.get_process_group_ranks(device_mesh["cp"].get_group()),
-                torch.cuda.Stream()
+                torch.cuda.Stream(),
             )
 
     dist.barrier()
@@ -149,18 +157,21 @@
 
     batch = get_dummy_data_bshd()
     input_ids, labels, position_ids = get_batch_on_this_cp_rank(
-        cu_seqlens_padded=batch['cu_seqlens_q'],
-        input_ids_padded=batch['input_ids'],
-        labels_padded=batch['labels'],
-        position_ids_padded=batch['position_ids'],
+        cu_seqlens_padded=batch["cu_seqlens_q"],
+        input_ids_padded=batch["input_ids"],
+        labels_padded=batch["labels"],
+        position_ids_padded=batch["position_ids"],
         cp_group=cp_group,
-        qvk_format="bshd"
+        qvk_format="bshd",
     )
-    batch['input_ids'] = input_ids
-    batch['labels'] = labels
-    batch['position_ids'] = position_ids
-    
-    batch = {k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v for k, v in batch.items()}    
+    batch["input_ids"] = input_ids
+    batch["labels"] = labels
+    batch["position_ids"] = position_ids
+
+    batch = {
+        k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v
+        for k, v in batch.items()
+    }
 
     with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
         output = model(batch)
@@ -177,19 +188,19 @@
     grads = collect_gradients(model, layer_patterns=target_layers, max_params=15)
 
     cp2_results = {
-        'logits': output.clone().detach().cpu(),
-        'loss': loss.clone().detach().cpu(),
-        'grad_norms': {name: grad.norm().item() for name, grad in grads.items()},
-        'grads': grads,
+        "logits": output.clone().detach().cpu(),
+        "loss": loss.clone().detach().cpu(),
+        "grad_norms": {name: grad.norm().item() for name, grad in grads.items()},
+        "grads": grads,
     }
 
-    torch.save(cp2_results, f'/tmp/bshd_cp2_rank_{dist_config.rank}_results.pt')
+    torch.save(cp2_results, f"/tmp/bshd_cp2_rank_{dist_config.rank}_results.pt")
     dist.barrier()
-    
+
     if dist_config.rank == 0:
         logger.info(f"CP=2 complete: rank0_loss={cp2_results['loss'].item():.6f}")
         # Create completion marker file to signal that all processing is done
-        with open('/tmp/bshd_complete.marker', 'w') as f:
-            f.write('completed')
-    
+        with open("/tmp/bshd_complete.marker", "w") as f:
+            f.write("completed")
+
     dist.destroy_process_group()

examples/pytorch/transformer/context_parallel_runner_thd.py

@@ -14,7 +14,9 @@
 
 from utils import get_dummy_data_thd, collect_gradients, DistributedConfig
 from model import SimpleConfig, SimpleThDModel
-from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import get_batch_on_this_cp_rank
+from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
+    get_batch_on_this_cp_rank,
+)
 
 
 random.seed(42)
@@ -23,12 +25,12 @@
 if torch.cuda.is_available():
     torch.cuda.manual_seed_all(42)
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 data = get_dummy_data_thd()
 
-DISTRIBUTED_MODE = 'RANK' in os.environ and 'WORLD_SIZE' in os.environ
+DISTRIBUTED_MODE = "RANK" in os.environ and "WORLD_SIZE" in os.environ
 
 device = torch.device("cuda:0")
 torch.cuda.set_device(0)
@@ -53,10 +55,13 @@
     torch.cuda.manual_seed(42)
 
 batch = get_dummy_data_thd()
-batch['input_ids'] = batch['input_ids_padded']
-batch['labels'] = batch['labels_padded']
-batch['position_ids'] = batch['position_ids_padded']
-batch = {k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+batch["input_ids"] = batch["input_ids_padded"]
+batch["labels"] = batch["labels_padded"]
+batch["position_ids"] = batch["position_ids_padded"]
+batch = {
+    k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v
+    for k, v in batch.items()
+}
 
 with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
     output_cp1 = model_cp1(batch)
@@ -65,69 +70,72 @@
 labels_flat = batch["labels"].view(-1)
 valid_mask = labels_flat != -100
 if valid_mask.any():
-    loss_cp1 = F.cross_entropy(logits_flat_cp1[valid_mask], labels_flat[valid_mask], reduction="mean")
+    loss_cp1 = F.cross_entropy(
+        logits_flat_cp1[valid_mask], labels_flat[valid_mask], reduction="mean"
+    )
 else:
     loss_cp1 = logits_flat_cp1.sum() * 0.0
 
 # Compute gradients for CP=1
 loss_cp1.backward()
 
 target_layers = [
-    'embedding',           # Embedding layers
-    'transformer_layers.0',      # First transformer layer
-    'transformer_layers.1',      # Second transformer layer  
-    'linear'              # Language model head
+    "embedding",  # Embedding layers
+    "transformer_layers.0",  # First transformer layer
+    "transformer_layers.1",  # Second transformer layer
+    "linear",  # Language model head
 ]
 
 grads_cp1 = collect_gradients(model_cp1, layer_patterns=target_layers, max_params=15)
 
 # logger.info(f"CP=1 collected {len(grads_cp1)} parameter gradients")
 
 initial_state_dict = {k: v.cpu().clone() for k, v in model_cp1.state_dict().items()}
-torch.save(initial_state_dict, '/tmp/thd_initial_model_state.pt')
+torch.save(initial_state_dict, "/tmp/thd_initial_model_state.pt")
 # logger.info("Model state saved for CP=2 reuse")
 
 cp1_results = {
-    'logits': output_cp1.clone().detach().cpu(),
-    'loss': loss_cp1.clone().detach().cpu(),
-    'grad_norms': {name: grad.norm().item() for name, grad in grads_cp1.items()},
-    'grads': grads_cp1,
+    "logits": output_cp1.clone().detach().cpu(),
+    "loss": loss_cp1.clone().detach().cpu(),
+    "grad_norms": {name: grad.norm().item() for name, grad in grads_cp1.items()},
+    "grads": grads_cp1,
 }
 
-torch.save(cp1_results, '/tmp/thd_cp1_results.pt')
-torch.save(data, '/tmp/thd_data.pt')
+torch.save(cp1_results, "/tmp/thd_cp1_results.pt")
+torch.save(data, "/tmp/thd_data.pt")
 
 # STEP 2: RUN CP=2 (CONTEXT PARALLELISM)
 # Skip CP=2 if not in distributed mode
 if not DISTRIBUTED_MODE:
-    logger.info("="*50)
+    logger.info("=" * 50)
     logger.info("SKIPPING CP=2 - Not running in distributed mode")
     logger.info("To test CP=2, run with: torchrun --nproc_per_node=2 run_context_parallel_thd.py")
-    logger.info("="*50)
+    logger.info("=" * 50)
     import sys
+
     sys.exit(1)
 
 # Run CP=2 in distributed mode
 if DISTRIBUTED_MODE:
-    logger.info("="*50)
+    logger.info("=" * 50)
     logger.info("RUNNING CP=2 (CONTEXT PARALLELISM)")
-    logger.info("="*50)
+    logger.info("=" * 50)
 
     cp_size = 2
     dist.init_process_group(backend="nccl")
     dist_config = DistributedConfig()
     torch.cuda.set_device(dist_config.local_rank)
 
     device_mesh = init_device_mesh(
-            "cuda",
-            mesh_shape=(1, cp_size, 1),
-            mesh_dim_names=("fsdp", "cp", "tp"),
-        )
+        "cuda",
+        mesh_shape=(1, cp_size, 1),
+        mesh_dim_names=("fsdp", "cp", "tp"),
+    )
     device = torch.device(f"cuda:{dist_config.local_rank}")
     model = SimpleThDModel(config)
 
     try:
-        initial_state_dict = torch.load('/tmp/thd_initial_model_state.pt', map_location='cpu')
+        initial_state_dict = torch.load("/tmp/thd_initial_model_state.pt", map_location="cpu")
         model.load_state_dict(initial_state_dict)
         # logger.info(f"Rank {dist_config.rank}: Model state loaded successfully")
     except Exception as e:
@@ -152,7 +160,7 @@
             transformer_layer.set_context_parallel_group(
                 cp_group,
                 torch.distributed.get_process_group_ranks(device_mesh["cp"].get_group()),
-                torch.cuda.Stream()
+                torch.cuda.Stream(),
             )
 
     dist.barrier()
@@ -162,20 +170,23 @@
         torch.cuda.manual_seed(42)
 
     batch = get_dummy_data_thd()
-    
+
     input_ids_padded, labels_padded, position_ids_padded = get_batch_on_this_cp_rank(
-        cu_seqlens_padded=batch['cu_seqlens_q_padded'],
-        input_ids_padded=batch['input_ids_padded'],
-        labels_padded=batch['labels_padded'],
-        position_ids_padded=batch['position_ids_padded'],
+        cu_seqlens_padded=batch["cu_seqlens_q_padded"],
+        input_ids_padded=batch["input_ids_padded"],
+        labels_padded=batch["labels_padded"],
+        position_ids_padded=batch["position_ids_padded"],
         cp_group=cp_group,
-        qvk_format="thd"
+        qvk_format="thd",
     )
 
-    batch['input_ids'] = input_ids_padded
-    batch['labels'] = labels_padded
-    batch['position_ids'] = position_ids_padded
-    batch = {k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+    batch["input_ids"] = input_ids_padded
+    batch["labels"] = labels_padded
+    batch["position_ids"] = position_ids_padded
+    batch = {
+        k: v.to(device, non_blocking=True).contiguous() if isinstance(v, torch.Tensor) else v
+        for k, v in batch.items()
+    }
 
     with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
         output = model(batch)
@@ -193,18 +204,18 @@
     grads = collect_gradients(model, layer_patterns=target_layers, max_params=15)
 
     cp2_results = {
-        'logits': output.clone().detach().cpu(),
-        'loss': loss.clone().detach().cpu(),
-        'grad_norms': {name: grad.norm().item() for name, grad in grads.items()},
-        'grads': grads,
+        "logits": output.clone().detach().cpu(),
+        "loss": loss.clone().detach().cpu(),
+        "grad_norms": {name: grad.norm().item() for name, grad in grads.items()},
+        "grads": grads,
     }
 
-    torch.save(cp2_results, f'/tmp/thd_cp2_rank_{dist_config.rank}_results.pt')
+    torch.save(cp2_results, f"/tmp/thd_cp2_rank_{dist_config.rank}_results.pt")
     dist.barrier()
-    
+
     # Create completion marker on rank 0 to signal all processing is done
     if dist_config.rank == 0:
-        with open('/tmp/thd_complete.marker', 'w') as f:
-            f.write('completed')
-    
+        with open("/tmp/thd_complete.marker", "w") as f:
+            f.write("completed")
+
     dist.destroy_process_group()