Fix isort indentation format

Author: yzhangcs

native_sparse_attention/modeling_nsa.py

@@ -11,8 +11,7 @@
 import torch.utils.checkpoint
 from einops import rearrange
 from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (BaseModelOutputWithPast,
-                                           CausalLMOutputWithPast)
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 from transformers.utils.deprecation import deprecate_kwarg

native_sparse_attention/ops/parallel.py

@@ -9,8 +9,7 @@
 import triton.language as tl
 from einops import rearrange
 
-from fla.ops.common.utils import (prepare_chunk_indices, prepare_chunk_offsets,
-                                  prepare_lens, prepare_token_indices)
+from fla.ops.common.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_lens, prepare_token_indices
 from fla.ops.utils import mean_pooling
 from fla.utils import autocast_custom_bwd, autocast_custom_fwd, contiguous
 from native_sparse_attention.ops.utils import _bitonic_merge

train.py

@@ -15,20 +15,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
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
 
 import native_sparse_attention  # noqa