Fix isort config and files

Add isort config based on the fms setup.
Fix imports where required. Include the root folder.

Signed-off-by: Andrea Frittoli <[email protected]>

.github/workflows/lint.yml

@@ -17,5 +17,5 @@ jobs:
           version: "~= 23.3.0"
       - uses: isort/isort-action@master
         with:
-          sort-paths: pretraining
+          sort-paths: .
           requirementsFiles: "requirements.txt" # We don't need extra test requirements for linting

.isort.cfg

@@ -0,0 +1,8 @@
+[settings]
+ensure_newline_before_comments = True
+force_grid_wrap = 0
+include_trailing_comma = True
+lines_after_imports = 2
+multi_line_output = 3
+use_parentheses = True
+profile = black

README.md

@@ -1 +1 @@
-# fms-pretrain
+# pretrain

pretraining/policies/ac_handler.py

@@ -7,6 +7,7 @@
     checkpoint_wrapper,
 )
 
+
 non_reentrant_wrapper = partial(
     checkpoint_wrapper,
     checkpoint_impl=CheckpointImpl.NO_REENTRANT,

pretraining/policies/mixed_precision.py

@@ -1,6 +1,7 @@
 import torch
 from torch.distributed.fsdp import MixedPrecision
 
+
 fpSixteen = MixedPrecision(
     param_dtype=torch.float16,
     reduce_dtype=torch.float16,

pretraining/utils/dataset_utils.py

@@ -4,15 +4,7 @@
 import os
 import random
 import sys
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Type,
-    Union,
-)
+from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 import pyarrow as pa
 import torch

pretraining/utils/train_utils.py

@@ -1,13 +1,16 @@
 import os
+
+
 try:
     import packaging.version
 except ImportError:
     from pkg_resources import packaging  # type: ignore
+
 import time
 
 import torch.cuda.nccl as nccl
 import torch.distributed as dist
-
+from packaging import version
 from torch.distributed.fsdp import ShardingStrategy
 
 from pretraining.policies import *

test-requirements.txt

@@ -12,6 +12,3 @@ mypy-extensions==1.0.0
 pyarrow-stubs==10.0.1.7
 types-requests==2.31.0.20240125
 types-setuptools==69.0.0.20240125
-
-# Install ibm-fms from the main branch for testing purposes
-ibm-fms @ git+https://github.com/foundation-model-stack/foundation-model-stack@main