fix format

Signed-off-by: yiliu30 <yi4.liu@intel.com>

Author: yiliu30

examples/autoround/ddp/ddp_qwen3_example.py

@@ -1,17 +1,16 @@
-#############################################################################
-# This script is adapted to use DDP functionality with AutoRound.
-# run this with `torchrun --nproc_per_node=2 ddp_qwen3_example.py`
-# or change nproc_per_node to your desired configuration
-#
-# Example usage:
-# torchrun --nproc_per_node=2 ddp_qwen3_example.py \
-#     --model Qwen/Qwen3-8B \
-#     --nsamples 128 \
-#     --iters 100 \
-#     --disable_torch_compile \
-#     --deterministic
-#############################################################################
-
+"""
+This script is adapted to use DDP functionality with AutoRound.
+run this with `torchrun --nproc_per_node=2 ddp_qwen3_example.py`
+or change nproc_per_node to your desired configuration
+
+Example usage:
+torchrun --nproc_per_node=2 ddp_qwen3_example.py \
+    --model Qwen/Qwen3-8B \
+    --nsamples 128 \
+    --iters 100 \
+    --disable_torch_compile \
+    --deterministic
+"""
 import argparse
 import os
 
@@ -93,12 +92,12 @@ def config_deterministic():
 
 
 # Get aligned calibration dataset.
-from auto_round.calib_dataset import get_dataset
+from auto_round.calib_dataset import get_dataset  # noqa: E402
 
 # Note: Make sure model are loaded before importing auto-round related code.
 # This requirement will be lifted once switching to new release of auto-round which
 # includes below fix:
-from llmcompressor.modifiers.autoround import AutoRoundModifier
+from llmcompressor.modifiers.autoround import AutoRoundModifier  # noqa: E402
 
 ds = get_dataset(
     tokenizer=tokenizer,
@@ -131,27 +130,28 @@ def config_deterministic():
 
 rank = dist.get_rank()
 logger.info(f"[Rank {rank}] Quantization completed")
-# Confirm generations of the quantized model look sane.
-logger.info("\n\n")
-logger.info("========== SAMPLE GENERATION ==============")
-dispatch_model(model)
-sample = tokenizer("Hello my name is", return_tensors="pt")
-sample = {key: value.to(model.device) for key, value in sample.items()}
-output = model.generate(**sample, max_new_tokens=100)
-logger.info(tokenizer.decode(output[0]))
-logger.info("==========================================\n\n")
-
-logger.info("Saving...")
-# Save to disk compressed.
-SAVE_DIR = (
-    model_id.rstrip("/").split("/")[-1]
-    + f"-{args.scheme}-AutoRound"
-    + f"-iters{args.iters}-nsamples{args.nsamples}"
-    + "-DDP"
-    + str(dist.get_world_size())
-)
-model.save_pretrained(SAVE_DIR, save_compressed=True)
-tokenizer.save_pretrained(SAVE_DIR)
-logger.info(f"Saved to {SAVE_DIR}")
+if rank == 0:
+    # Confirm generations of the quantized model look sane.
+    logger.info("\n\n")
+    logger.info("========== SAMPLE GENERATION ==============")
+    dispatch_model(model)
+    sample = tokenizer("Hello my name is", return_tensors="pt")
+    sample = {key: value.to(model.device) for key, value in sample.items()}
+    output = model.generate(**sample, max_new_tokens=100)
+    logger.info(tokenizer.decode(output[0]))
+    logger.info("==========================================\n\n")
+
+    logger.info("Saving...")
+    # Save to disk compressed.
+    SAVE_DIR = (
+        model_id.rstrip("/").split("/")[-1]
+        + f"-{args.scheme}-AutoRound"
+        + f"-iters{args.iters}-nsamples{args.nsamples}"
+        + "-DDP"
+        + str(dist.get_world_size())
+    )
+    model.save_pretrained(SAVE_DIR, save_compressed=True)
+    tokenizer.save_pretrained(SAVE_DIR)
+    logger.info(f"Saved to {SAVE_DIR}")
 
 dist.destroy_process_group()