Fix isort indentation

Author: yzhangcs

3rdparty/flash-linear-attention

@@ -12,19 +12,18 @@
 import torch
 import torch.nn as nn
 from torch.distributed import DeviceMesh
-from torch.distributed._composable.fsdp import (CPUOffloadPolicy,
-                                                MixedPrecisionPolicy,
-                                                fully_shard)
+from torch.distributed._composable.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
 from torch.distributed._composable.replicate import replicate
 from torch.distributed._tensor import Replicate, Shard
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import \
-    checkpoint_wrapper as ptd_checkpoint_wrapper
-from torch.distributed.tensor.parallel import (ColwiseParallel,
-                                               PrepareModuleInput,
-                                               PrepareModuleOutput,
-                                               RowwiseParallel,
-                                               SequenceParallel,
-                                               parallelize_module)
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper as ptd_checkpoint_wrapper
+from torch.distributed.tensor.parallel import (
+    ColwiseParallel,
+    PrepareModuleInput,
+    PrepareModuleOutput,
+    RowwiseParallel,
+    SequenceParallel,
+    parallelize_module
+)
 
 from fla.modules.fused_linear_cross_entropy import LinearLossParallel
 from fla.modules.mlp import SwiGLULinearParallel
@@ -126,8 +125,10 @@ def __init__(
         # TODO(vkuzo): add the items below to __init__.py of torchao.float8 and import from there
         try:
             from torchao.float8.float8_tensor_parallel import (
-                Float8ColwiseParallel, Float8RowwiseParallel,
-                PrepareFloat8ModuleInput)
+                Float8ColwiseParallel,
+                Float8RowwiseParallel,
+                PrepareFloat8ModuleInput
+            )
         except ImportError:
             Float8ColwiseParallel = None
             Float8RowwiseParallel = None
@@ -268,8 +269,7 @@ def apply_tp(
             )
 
     if enable_async_tp:
-        from torch.distributed._symmetric_memory import \
-            enable_symm_mem_for_group
+        from torch.distributed._symmetric_memory import enable_symm_mem_for_group
 
         torch._inductor.config._micro_pipeline_tp = True
         enable_symm_mem_for_group(tp_mesh.get_group().group_name)
@@ -312,8 +312,7 @@ def _apply_ac_to_block(module: nn.Module, ac_config):
             f"Valid options: 'op' or a positive int representing layer frequency"
         )
     if use_op_sac:
-        from torch.utils.checkpoint import (
-            CheckpointPolicy, create_selective_checkpoint_contexts)
+        from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
 
         def _get_custom_policy(meta):
             def _custom_policy(ctx, func, *args, **kwargs):

flame/models/parallelize_fla.py

@@ -13,18 +13,13 @@
 import torch.nn as nn
 from torch.distributed import DeviceMesh
 from torch.distributed.pipelining import PipelineStage
-from torch.distributed.pipelining.schedules import (ScheduleZBVZeroBubble,
-                                                    _PipelineSchedule,
-                                                    get_schedule_class)
+from torch.distributed.pipelining.schedules import ScheduleZBVZeroBubble, _PipelineSchedule, get_schedule_class
 from transformers import PretrainedConfig
 
-from flame.models.parallelize_fla import (get_blocks, get_components_name,
-                                          get_model)
+from flame.models.parallelize_fla import get_blocks, get_components_name, get_model
 from torchtitan.config_manager import JobConfig
 from torchtitan.distributed.parallel_dims import ParallelDims
-from torchtitan.distributed.pipeline import (build_pipeline_schedule,
-                                             generate_split_points,
-                                             stage_ids_this_rank)
+from torchtitan.distributed.pipeline import build_pipeline_schedule, generate_split_points, stage_ids_this_rank
 from torchtitan.tools.logging import logger
 
 DeviceType = Union[int, str, torch.device]

flame/models/pipeline_fla.py

@@ -37,3 +37,7 @@ Homepage = "https://github.com/fla-org/flame"
 [build-system]
 requires = ["setuptools>=45", "wheel", "ninja", "torch"]
 build-backend = "setuptools.build_meta"
+
+[tool.isort]
+line_length = 127
+multi_line_output = 3

pyproject.toml

@@ -26,20 +26,15 @@
 from torchtitan.components.checkpoint import CheckpointManager
 from torchtitan.components.ft import FTParallelDims, init_ft_manager
 from torchtitan.components.loss import cross_entropy_loss
-from torchtitan.components.metrics import (_build_metric_logger,
-                                           build_device_memory_monitor,
-                                           ensure_pp_loss_visible)
-from torchtitan.components.optimizer import (build_lr_schedulers,
-                                             build_optimizers)
+from torchtitan.components.metrics import _build_metric_logger, build_device_memory_monitor, ensure_pp_loss_visible
+from torchtitan.components.optimizer import build_lr_schedulers, build_optimizers
 from torchtitan.distributed import ParallelDims
 from torchtitan.distributed import utils as dist_utils
 from torchtitan.protocols.model_converter import build_model_converters
-from torchtitan.protocols.train_spec import (TrainSpec, get_train_spec,
-                                             register_train_spec)
+from torchtitan.protocols.train_spec import TrainSpec, get_train_spec, register_train_spec
 from torchtitan.tools import utils
 from torchtitan.tools.logging import init_logger, logger
-from torchtitan.tools.profiling import (maybe_enable_memory_snapshot,
-                                        maybe_enable_profiling)
+from torchtitan.tools.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling
 
 register_train_spec(
     TrainSpec(