remove unused import

dcu_megatron/core/models/gpt/fine_grained_schedule.py

@@ -773,7 +773,7 @@ def build_model_chunk_schedule_plan(
     state.attention_mask = attention_mask
     state.decoder_input = decoder_input
     state.labels = labels
-    state.inference_context =inference_context
+    state.inference_context = inference_context
     state.packed_seq_params = packed_seq_params
     state.extra_block_kwargs = extra_block_kwargs
     state.runtime_gather_output = runtime_gather_output

dcu_megatron/core/pipeline_parallel/combined_1f1b.py

 import contextlib
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import Any, List, Tuple, Union
+from typing import List, Union
 
 import torch
 from torch import Tensor
@@ -16,10 +16,6 @@ from megatron.core.transformer.multi_token_prediction import MTPLossAutoScaler
 from megatron.core.utils import get_attr_wrapped_model, make_viewless_tensor
 
 
-# Types
-Shape = Union[List[int], torch.Size]
-
-
 def make_viewless(e):
     """make_viewless util func"""
     e = make_viewless_tensor(inp=e, requires_grad=e.requires_grad, keep_graph=True)
@@ -351,7 +347,7 @@ def forward_backward_step(
         Tensor or list[Tensor]: The output object(s) from the forward step.
         Tensor: The number of tokens.
     """
-    from .schedules import set_current_microbatch
+    from megatron.core.pipeline_parallel.schedules import set_current_microbatch
 
     if config.timers is not None:
         config.timers('forward-compute', log_level=2).start()

dcu_megatron/core/pipeline_parallel/schedules.py

 import contextlib
-from typing import Callable, Iterator, List, Optional, Union
+from typing import Iterator, List, Union
 
 import torch
 
@@ -7,10 +7,8 @@ from megatron.training import get_args
 from megatron.core import parallel_state
 from megatron.core.enums import ModelType
 from megatron.core.pipeline_parallel import p2p_communication
-from megatron.core.pipeline_parallel.schedules import set_current_microbatch
 from megatron.core.transformer.cuda_graphs import create_cudagraphs
 from megatron.core.utils import (
-    get_attr_wrapped_model,
     get_model_config,
     get_model_type,
     get_model_xattn,

dcu_megatron/core/transformer/transformer_layer.py

-from functools import partial
 from typing import Any, Optional
 
-import torch
 from torch import Tensor
 
 from megatron.core import tensor_parallel