diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml
new file mode 100644
index 000000000..9b54a269f
--- /dev/null
+++ b/.github/workflows/quality.yml
@@ -0,0 +1,17 @@
+name: Quality Check
+
+on: [pull_request]
+
+jobs:
+  quality:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.6
+    - name: Install Python dependencies
+      run: pip install -e .[quality]
+    - name: Run Quality check
+      run: make quality
\ No newline at end of file
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 000000000..aa04063d8
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,17 @@
+name: Run Tests
+
+on: [pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.6
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.6
+    - name: Install Python dependencies
+      run: pip install -e .[test]
+    - name: Run Tests
+      run: make test
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 7a7a6a9cd..0fdc5e785 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,21 @@
 
 extras = {}
 extras["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]
-extras["docs"] = ["recommonmark", "sphinx==3.2.1", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
+extras["docs"] = [
+    "recommonmark",
+    "sphinx==3.2.1",
+    "sphinx-markdown-tables",
+    "sphinx-rtd-theme==0.4.3",
+    "sphinx-copybutton",
+]
+extras["test"] = [
+    "pytest",
+    "pytest-xdist",
+]
+
+extras["sagemaker"] = [
+    "sagemaker",  # boto3 is a required package in sagemaker
+]
 
 setup(
     name="accelerate",
@@ -32,11 +46,13 @@
     url="https://github.com/huggingface/accelerate",
     package_dir={"": "src"},
     packages=find_packages("src"),
-    entry_points={"console_scripts": [
-        "accelerate=accelerate.commands.accelerate_cli:main",
-        "accelerate-config=accelerate.commands.config:main",
-        "accelerate-launch=accelerate.commands.launch:main",
-    ]},
+    entry_points={
+        "console_scripts": [
+            "accelerate=accelerate.commands.accelerate_cli:main",
+            "accelerate-config=accelerate.commands.config:main",
+            "accelerate-launch=accelerate.commands.launch:main",
+        ]
+    },
     python_requires=">=3.6.0",
     install_requires=["torch>=1.4.0"],
     extras_require=extras,
diff --git a/src/accelerate/commands/config.py b/src/accelerate/commands/config.py
deleted file mode 100644
index c32905932..000000000
--- a/src/accelerate/commands/config.py
+++ /dev/null
@@ -1,218 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2021 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import argparse
-import json
-import os
-from dataclasses import dataclass
-from enum import Enum
-from typing import Optional
-
-import yaml
-from accelerate.state import DistributedType
-
-
-hf_cache_home = os.path.expanduser(
-    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
-)
-cache_dir = os.path.join(hf_cache_home, "accelerate")
-default_json_config_file = os.path.join(cache_dir, "default_config.json")
-default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
-
-# For backward compatibility: the default config is the json one if it's the only existing file.
-if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
-    default_config_file = default_yaml_config_file
-else:
-    default_config_file = default_json_config_file
-
-
-@dataclass
-class LaunchConfig:
-    distributed_type: DistributedType
-    num_processes: int
-    fp16: bool
-    machine_rank: int = 0
-    num_machines: int = 1
-    main_process_ip: Optional[str] = None
-    main_process_port: Optional[int] = None
-    main_training_function: str = "main"
-
-    def __post_init__(self):
-        if isinstance(self.distributed_type, str):
-            self.distributed_type = DistributedType(self.distributed_type)
-
-    def to_dict(self):
-        result = self.__dict__
-        # For serialization, it's best to convert Enums to strings (or their underlying value type).
-        for key, value in result.items():
-            if isinstance(value, Enum):
-                result[key] = value.value
-        return result
-
-    @classmethod
-    def from_json_file(cls, json_file=None):
-        json_file = default_json_config_file if json_file is None else json_file
-        with open(json_file, "r", encoding="utf-8") as f:
-            return cls(**json.load(f))
-
-    def to_json_file(self, json_file):
-        with open(json_file, "w", encoding="utf-8") as f:
-            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
-            f.write(content)
-
-    @classmethod
-    def from_yaml_file(cls, yaml_file=None):
-        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
-        with open(yaml_file, "r", encoding="utf-8") as f:
-            return cls(**yaml.safe_load(f))
-
-    def to_yaml_file(self, yaml_file):
-        with open(yaml_file, "w", encoding="utf-8") as f:
-            yaml.safe_dump(self.to_dict(), f)
-
-
-def config_command_parser(subparsers=None):
-    if subparsers is not None:
-        parser = subparsers.add_parser("config")
-    else:
-        parser = argparse.ArgumentParser("Accelerate config command")
-
-    parser.add_argument(
-        "--config_file",
-        default=None,
-        help=(
-            "The path to use to store the config file. Will default to a file named default_config.json in the cache "
-            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
-            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
-            "with 'huggingface'."
-        ),
-    )
-
-    if subparsers is not None:
-        parser.set_defaults(func=config_command)
-    return parser
-
-
-def _ask_field(input_text, convert_value=None, default=None, error_message=None):
-    ask_again = True
-    while ask_again:
-        result = input(input_text)
-        try:
-            if default is not None and len(result) == 0:
-                return default
-            return convert_value(result) if convert_value is not None else result
-        except:
-            if error_message is not None:
-                print(error_message)
-
-
-def get_user_input():
-    def _convert_distributed_mode(value):
-        value = int(value)
-        return DistributedType(["NO", "MULTI_GPU", "TPU"][value])
-
-    def _convert_yes_no_to_bool(value):
-        return {"yes": True, "no": False}[value.lower()]
-
-    distributed_type = _ask_field(
-        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
-        _convert_distributed_mode,
-        error_message="Please enter 0, 1 or 2.",
-    )
-
-    machine_rank = 0
-    num_machines = 1
-    main_process_ip = None
-    main_process_port = None
-    if distributed_type == DistributedType.MULTI_GPU:
-        num_machines = _ask_field(
-            "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
-            lambda x: int(x),
-            default=1,
-        )
-        if num_machines > 1:
-            machine_rank = _ask_field(
-                "What is the rank of this machine (from 0 to the number of machines - 1 )? [0]: ",
-                lambda x: int(x),
-                default=0,
-            )
-            main_process_ip = _ask_field(
-                "What is the IP address of the machine that will host the main process? ",
-            )
-            main_process_ip = _ask_field(
-                "What is the port you will use to communicate with the main process? ",
-                lambda x: int(x),
-            )
-    if distributed_type == DistributedType.TPU:
-        main_training_function = _ask_field(
-            "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
-            default="main",
-        )
-    else:
-        main_training_function = "main"
-
-    num_processes = _ask_field(
-        "How many processes in total will you use? [1]: ",
-        lambda x: int(x),
-        default=1,
-        error_message="Please enter an integer.",
-    )
-
-    if distributed_type != DistributedType.TPU:
-        fp16 = _ask_field(
-            "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
-            _convert_yes_no_to_bool,
-            default=False,
-            error_message="Please enter yes or no.",
-        )
-    else:
-        fp16 = False
-
-    return LaunchConfig(
-        distributed_type=distributed_type,
-        num_processes=num_processes,
-        fp16=fp16,
-        machine_rank=machine_rank,
-        num_machines=num_machines,
-        main_process_ip=main_process_ip,
-        main_process_port=main_process_port,
-        main_training_function=main_training_function,
-    )
-
-
-def config_command(args):
-    config = get_user_input()
-    if args.config_file is not None:
-        config_file = args.config_file
-    else:
-        if not os.path.isdir(cache_dir):
-            os.makedirs(cache_dir)
-        config_file = default_yaml_config_file
-
-    if config_file.endswith(".json"):
-        config.to_json_file(config_file)
-    else:
-        config.to_yaml_file(config_file)
-
-
-def main():
-    parser = config_command_parser()
-    args = parser.parse_args()
-    config_command(args)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/src/accelerate/commands/config/__init__.py b/src/accelerate/commands/config/__init__.py
new file mode 100644
index 000000000..48ea96ab1
--- /dev/null
+++ b/src/accelerate/commands/config/__init__.py
@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from accelerate.state import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
+from .config_utils import _ask_field, _convert_compute_environment
+from .sagemaker import get_sagemaker_input
+
+
+def get_user_input():
+    compute_environment = _ask_field(
+        "In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): ",
+        _convert_compute_environment,
+        error_message="Please enter 0 or 1",
+    )
+    if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+        config = get_sagemaker_input()
+    else:
+        config = get_cluster_input()
+    return config
+
+
+def config_command_parser(subparsers=None):
+    if subparsers is not None:
+        parser = subparsers.add_parser("config")
+    else:
+        parser = argparse.ArgumentParser("Accelerate config command")
+
+    parser.add_argument(
+        "--config_file",
+        default=None,
+        help=(
+            "The path to use to store the config file. Will default to a file named default_config.json in the cache "
+            "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
+            "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
+            "with 'huggingface'."
+        ),
+    )
+
+    if subparsers is not None:
+        parser.set_defaults(func=config_command)
+    return parser
+
+
+def config_command(args):
+    config = get_user_input()
+    if args.config_file is not None:
+        config_file = args.config_file
+    else:
+        if not os.path.isdir(cache_dir):
+            os.makedirs(cache_dir)
+        config_file = default_yaml_config_file
+
+    if config_file.endswith(".json"):
+        config.to_json_file(config_file)
+    else:
+        config.to_yaml_file(config_file)
+
+
+def main():
+    parser = config_command_parser()
+    args = parser.parse_args()
+    config_command(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/accelerate/commands/config/cluster.py b/src/accelerate/commands/config/cluster.py
new file mode 100644
index 000000000..a9e412f7d
--- /dev/null
+++ b/src/accelerate/commands/config/cluster.py
@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from accelerate.state import ComputeEnvironment, DistributedType
+
+from .config_args import ClusterConfig
+from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
+
+
+def get_cluster_input():
+    distributed_type = _ask_field(
+        "Which type of machine are you using? ([0] No distributed training, [1] multi-GPU, [2] TPU): ",
+        _convert_distributed_mode,
+        error_message="Please enter 0, 1 or 2.",
+    )
+
+    machine_rank = 0
+    num_machines = 1
+    main_process_ip = None
+    main_process_port = None
+    if distributed_type == DistributedType.MULTI_GPU:
+        num_machines = _ask_field(
+            "How many different machines will you use (use more than 1 for multi-node training)? [1]: ",
+            lambda x: int(x),
+            default=1,
+        )
+        if num_machines > 1:
+            machine_rank = _ask_field(
+                "What is the rank of this machine (from 0 to the number of machines - 1 )? [0]: ",
+                lambda x: int(x),
+                default=0,
+            )
+            main_process_ip = _ask_field(
+                "What is the IP address of the machine that will host the main process? ",
+            )
+            main_process_ip = _ask_field(
+                "What is the port you will use to communicate with the main process? ",
+                lambda x: int(x),
+            )
+    if distributed_type == DistributedType.TPU:
+        main_training_function = _ask_field(
+            "What is the name of the function in your script that should be launched in all parallel scripts? [main]: ",
+            default="main",
+        )
+    else:
+        main_training_function = "main"
+
+        num_processes = _ask_field(
+            "How many processes in total will you use? [1]: ",
+            lambda x: int(x),
+            default=1,
+            error_message="Please enter an integer.",
+        )
+
+        if distributed_type != DistributedType.TPU:
+            fp16 = _ask_field(
+                "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
+                _convert_yes_no_to_bool,
+                default=False,
+                error_message="Please enter yes or no.",
+            )
+        else:
+            fp16 = False
+
+    return ClusterConfig(
+        compute_environment=ComputeEnvironment.LOCAL_MACHINE,
+        distributed_type=distributed_type,
+        num_processes=num_processes,
+        fp16=fp16,
+        machine_rank=machine_rank,
+        num_machines=num_machines,
+        main_process_ip=main_process_ip,
+        main_process_port=main_process_port,
+        main_training_function=main_training_function,
+    )
diff --git a/src/accelerate/commands/config/config_args.py b/src/accelerate/commands/config/config_args.py
new file mode 100644
index 000000000..5b0dd00ba
--- /dev/null
+++ b/src/accelerate/commands/config/config_args.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.LOCAL_MACHINE:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.LOCAL_MACHINE:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    fp16: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, "r", encoding="utf-8") as f:
+            return cls(**json.load(f))
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.compute_environment, str):
+            self.compute_environment = ComputeEnvironment(self.compute_environment)
+        if isinstance(self.distributed_type, str):
+            self.distributed_type = DistributedType(self.distributed_type)
+
+
+@dataclass
+class ClusterConfig(BaseConfig):
+    num_processes: int
+    machine_rank: int = 0
+    num_machines: int = 1
+    main_process_ip: Optional[str] = None
+    main_process_port: Optional[int] = None
+    main_training_function: str = "main"
+
+
+@dataclass
+class SageMakerConfig(BaseConfig):
+    ec2_instance_type: str
+    iam_role_name: str
+    profile: Optional[str] = None
+    region: str = "us-east-1"
+    num_machines: int = 1
+    base_job_name: str = f"accelerate-sagemaker-{num_machines}"
+    pytroch_version: str = "1.6"
+    transformers_version: str = "4.4"
diff --git a/src/accelerate/commands/config/config_utils.py b/src/accelerate/commands/config/config_utils.py
new file mode 100644
index 000000000..dd14965f7
--- /dev/null
+++ b/src/accelerate/commands/config/config_utils.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+def _ask_field(input_text, convert_value=None, default=None, error_message=None):
+    ask_again = True
+    while ask_again:
+        result = input(input_text)
+        try:
+            if default is not None and len(result) == 0:
+                return default
+            return convert_value(result) if convert_value is not None else result
+        except:
+            if error_message is not None:
+                print(error_message)
+
+
+def _convert_compute_environment(value):
+    value = int(value)
+    return ComputeEnvironment(["LOCAL_MACHINE", "AMAZON_SAGEMAKER"][value])
+
+
+def _convert_distributed_mode(value):
+    value = int(value)
+    return DistributedType(["NO", "MULTI_GPU", "TPU"][value])
+
+
+def _convert_sagemaker_distributed_mode(value):
+    value = int(value)
+    return SageMakerDistributedType(["NO", "DATA_PARALLEL", "MODEL_PARALLEL"][value])
+
+
+def _convert_yes_no_to_bool(value):
+    return {"yes": True, "no": False}[value.lower()]
diff --git a/src/accelerate/commands/config/sagemaker.py b/src/accelerate/commands/config/sagemaker.py
new file mode 100644
index 000000000..25fd3e6cb
--- /dev/null
+++ b/src/accelerate/commands/config/sagemaker.py
@@ -0,0 +1,157 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import json
+import os
+
+from accelerate.state import ComputeEnvironment, SageMakerDistributedType
+from accelerate.utils import is_boto3_available
+
+from .config_args import SageMakerConfig
+from .config_utils import _ask_field, _convert_sagemaker_distributed_mode, _convert_yes_no_to_bool
+
+
+if is_boto3_available():
+    import boto3  # noqa: F401
+
+
+def _create_iam_role_for_sagemaker(role_name):
+    iam_client = boto3.client("iam")
+
+    sagemaker_trust_policy = {
+        "Version": "2012-10-17",
+        "Statement": [
+            {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"}
+        ],
+    }
+    try:
+        # create the role, associated with the chosen trust policy
+        iam_client.create_role(
+            RoleName=role_name, AssumeRolePolicyDocument=json.dumps(sagemaker_trust_policy, indent=2)
+        )
+        policy_document = {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Action": [
+                        "sagemaker:*",
+                        "ecr:GetDownloadUrlForLayer",
+                        "ecr:BatchGetImage",
+                        "ecr:BatchCheckLayerAvailability",
+                        "ecr:GetAuthorizationToken",
+                        "cloudwatch:PutMetricData",
+                        "cloudwatch:GetMetricData",
+                        "cloudwatch:GetMetricStatistics",
+                        "cloudwatch:ListMetrics",
+                        "logs:CreateLogGroup",
+                        "logs:CreateLogStream",
+                        "logs:DescribeLogStreams",
+                        "logs:PutLogEvents",
+                        "logs:GetLogEvents",
+                        "s3:CreateBucket",
+                        "s3:ListBucket",
+                        "s3:GetBucketLocation",
+                        "s3:GetObject",
+                        "s3:PutObject",
+                    ],
+                    "Resource": "*",
+                }
+            ],
+        }
+        # attach policy to role
+        iam_client.put_role_policy(
+            RoleName=role_name,
+            PolicyName=f"{role_name}_policy_permission",
+            PolicyDocument=json.dumps(policy_document, indent=2),
+        )
+    except iam_client.exceptions.EntityAlreadyExistsException:
+        print(f"role {role_name} already exists. Using existing one")
+
+
+def _get_iam_role_arn(role_name):
+    iam_client = boto3.client("iam")
+    return iam_client.get_role(RoleName=role_name)["Role"]["Arn"]
+
+
+def get_sagemaker_input():
+    credentials_configuration = _ask_field(
+        "How do you want to authorize? ([0] AWS Profile, [1] Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)): ",
+        lambda x: int(x),
+    )
+    aws_profile = None
+    if credentials_configuration == 0:
+        aws_profile = _ask_field("Enter your AWS Profile name: [default] ", default="default")
+        os.environ["AWS_PROFILE"] = aws_profile
+    else:
+        print(
+            "Accelerate will expose AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables on the machine"
+        )
+        aws_access_key_id = _ask_field("AWS Access Key ID: ")
+        os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
+
+        aws_secret_access_key = _ask_field("AWS Secret Access Key: ")
+        os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
+
+    aws_region = _ask_field("Enter your AWS Region: [us-east-1]", default="us-east-1")
+    os.environ["AWS_DEFAULT_REGION"] = aws_region
+
+    role_management = _ask_field(
+        "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs? ([0] provide IAM Role name, [1] create new IAM role using credentials: ",
+        lambda x: int(x),
+    )
+    if role_management == 0:
+        iam_role_name = _ask_field("Enter your IAM role name: ")
+    else:
+        iam_role_name = "accelerate_sagemaker_execution_role"
+        print(f'Accelerate will create an iam role "{iam_role_name}" using the provided credentials')
+        _create_iam_role_for_sagemaker(iam_role_name)
+
+    distributed_type = _ask_field(
+        "Which type of machine are you using? ([0] No distributed training, [1] data parallelism, [2] model parallelism): ",
+        _convert_sagemaker_distributed_mode,
+        error_message="Please enter 0, 1 or 2",
+    )
+
+    # using the best two instances for single-gpu training or multi-gpu -> can turn into question to make it more diverse
+    ec2_instance_type = "ml.p3.2xlarge" if distributed_type == SageMakerDistributedType.NO else "ml.p3dn.24xlarge"
+    num_machines = 1
+    if (
+        distributed_type == SageMakerDistributedType.DATA_PARALLEL
+        or distributed_type == SageMakerDistributedType.MODEL_PARALLEL
+    ):
+        raise NotImplementedError("Model or Data Parallelism is not implemented yet. We are working on it")
+        num_machines = _ask_field(
+            "How many machines do you want use? [2]: ",
+            lambda x: int(x),
+            default=2,
+        )
+    fp16 = _ask_field(
+        "Do you wish to use FP16 (mixed precision)? [yes/NO]: ",
+        _convert_yes_no_to_bool,
+        default=False,
+        error_message="Please enter yes or no.",
+    )
+
+    return SageMakerConfig(
+        compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,
+        distributed_type=distributed_type,
+        ec2_instance_type=ec2_instance_type,
+        profile=aws_profile,
+        region=aws_region,
+        iam_role_name=iam_role_name,
+        fp16=fp16,
+        num_machines=num_machines,
+    )
diff --git a/src/accelerate/commands/launch.py b/src/accelerate/commands/launch.py
index 061c289ea..b8e1e8156 100644
--- a/src/accelerate/commands/launch.py
+++ b/src/accelerate/commands/launch.py
@@ -23,8 +23,8 @@
 from pathlib import Path
 from typing import Optional
 
-from accelerate.commands.config import LaunchConfig, default_config_file
-from accelerate.state import DistributedType
+from accelerate.commands.config import default_config_file, load_config_from_file
+from accelerate.state import ComputeEnvironment, DistributedType
 
 
 class _AddOneArg:
@@ -171,6 +171,12 @@ def tpu_launcher(args):
         xmp.spawn(main_function, args=(), nprocs=args.num_processes)
 
 
+def sagemaker_launcher(sagemaker_config, args):
+    raise NotImplementedError(
+        "Support for starting SageMaker training is not yet implemented. But you can create configs for it."
+    )
+
+
 def launch_command(args):
     # Sanity checks
     if args.multi_gpu and args.tpu:
@@ -178,15 +184,15 @@ def launch_command(args):
 
     # Get the default from the config file.
     if args.config_file is not None or os.path.isfile(default_config_file) and not args.cpu:
-        defaults = LaunchConfig.from_json_file(json_file=args.config_file)
+        defaults = load_config_from_file(args.config_file)
         if not args.multi_gpu and not args.tpu:
             args.multi_gpu = defaults.distributed_type == DistributedType.MULTI_GPU
             args.tpu = defaults.distributed_type == DistributedType.TPU
-        if args.num_processes is None:
+        if args.num_processes is None and defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             args.num_processes = defaults.num_processes
         if not args.fp16:
             args.fp16 = defaults.fp16
-        if args.main_training_function is None:
+        if args.main_training_function is None and defaults.compute_environment == ComputeEnvironment.LOCAL_MACHINE:
             args.main_training_function = defaults.main_training_function
     else:
         if args.num_processes is None:
@@ -197,6 +203,8 @@ def launch_command(args):
         multi_gpu_launcher(args)
     elif args.tpu and not args.cpu:
         tpu_launcher(args)
+    elif defaults.compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
+        sagemaker_launcher(defaults, args)
     else:
         simple_launcher(args)
 
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
index 5f150add9..55831be0f 100644
--- a/src/accelerate/state.py
+++ b/src/accelerate/state.py
@@ -53,6 +53,38 @@ class DistributedType(str, Enum):
     TPU = "TPU"
 
 
+class SageMakerDistributedType(str, Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **LOCAL_MACHINE** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    LOCAL_MACHINE = "LOCAL_MACHINE"
+    AMAZON_SAGEMAKER = "AMAZON_SAGEMAKER"
+
+
 # Inspired by Alex Martelli's 'Borg'.
 class AcceleratorState:
     """
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 61a39fc95..ca2d558b0 100644
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import importlib
 import random
 from enum import Enum
 from typing import List, Optional, Union
@@ -26,6 +27,10 @@
     import torch_xla.core.xla_model as xm
 
 
+def is_boto3_available():
+    return importlib.util.find_spec("boto3") is not None
+
+
 class RNGType(Enum):
     TORCH = "torch"
     CUDA = "cuda"


Review comments: body: decoupled input gathering from main config `get_user_input()` into the separate options
diff_hunk: @@ -0,0 +1,88 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from accelerate.state import ComputeEnvironment, DistributedType
+
+from .config_args import ClusterConfig
+from .config_utils import _ask_field, _convert_distributed_mode, _convert_yes_no_to_bool
+
+
+def get_cluster_input():
body: Determines which config flow should be used. 
diff_hunk: @@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from accelerate.state import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_yaml_config_file
+from .config_utils import _ask_field, _convert_compute_environment
+from .sagemaker import get_sagemaker_input
+
+
+def get_user_input():
+    compute_environment = _ask_field(
body: create a separate file to merge all configuration `dataclasses`. We can also think about moving them so `state.py`
diff_hunk: @@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
body: general loading function for determining if config file is `json` or `yaml` and which configuration is used for loading. Open for suggestions on how to identify the `config_class` easier. 
diff_hunk: @@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
body: base class contains overlapping properties and methods. 
diff_hunk: @@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
body: kept `num_machines` in specialist class due to `dataclass` -> `Dataclass fields without default value cannot appear after data fields with default values`. 
diff_hunk: @@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    fp16: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, "r", encoding="utf-8") as f:
+            return cls(**json.load(f))
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.distributed_type, str):
+            self.compute_environment = ComputeEnvironment(self.compute_environment)
+        if isinstance(self.distributed_type, str):
+            self.distributed_type = DistributedType(self.distributed_type)
+
+
+@dataclass
+class ClusterConfig(BaseConfig):
+    num_processes: int
+    machine_rank: int = 0
+    num_machines: int = 1
body: open for better naming
diff_hunk: @@ -53,6 +53,38 @@ class DistributedType(str, Enum):
     TPU = "TPU"
 
 
+class SageMakerDistributedType(str, Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **CUSTOM_CLUSTER** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    CUSTOM_CLUSTER = "CUSTOM_CLUSTER"
body: This is a bit weird so would remove one .config by adding `config_command_parser` in the intermediate init.
diff_hunk: @@ -16,7 +16,7 @@
 
 from argparse import ArgumentParser
 
-from accelerate.commands.config import config_command_parser
+from accelerate.commands.config.config import config_command_parser
body: Custom Cluster does sound a bit too complicated for the base config.
```suggestion
        "In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): ",
```
diff_hunk: @@ -0,0 +1,85 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+from accelerate.state import ComputeEnvironment
+
+from .cluster import get_cluster_input
+from .config_args import cache_dir, default_yaml_config_file
+from .config_utils import _ask_field, _convert_compute_environment
+from .sagemaker import get_sagemaker_input
+
+
+def get_user_input():
+    compute_environment = _ask_field(
+        "In which compute environment are you running? ([0] Custom Cluster, [1] AWS (Amazon SageMaker)): ",
body: ```suggestion
        if isinstance(self.compute_environment, str):
```
diff_hunk: @@ -0,0 +1,119 @@
+#!/usr/bin/env python
+
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Union
+
+import yaml
+from accelerate.state import ComputeEnvironment, DistributedType, SageMakerDistributedType
+
+
+hf_cache_home = os.path.expanduser(
+    os.getenv("HF_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "huggingface"))
+)
+cache_dir = os.path.join(hf_cache_home, "accelerate")
+default_json_config_file = os.path.join(cache_dir, "default_config.json")
+default_yaml_config_file = os.path.join(cache_dir, "default_config.yaml")
+
+# For backward compatibility: the default config is the json one if it's the only existing file.
+if os.path.isfile(default_yaml_config_file) or not os.path.isfile(default_json_config_file):
+    default_config_file = default_yaml_config_file
+else:
+    default_config_file = default_json_config_file
+
+
+def load_config_from_file(config_file):
+    config_file = config_file if config_file is not None else default_config_file
+    with open(config_file, "r", encoding="utf-8") as f:
+        if config_file.endswith(".json"):
+            if json.load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_json_file(json_file=config_file)
+        else:
+            if yaml.safe_load(f)["compute_environment"] == ComputeEnvironment.CUSTOM_CLUSTER:
+                config_class = ClusterConfig
+            else:
+                config_class = SageMakerConfig
+            return config_class.from_yaml_file(yaml_file=config_file)
+
+
+@dataclass
+class BaseConfig:
+    compute_environment: ComputeEnvironment
+    distributed_type: Union[DistributedType, SageMakerDistributedType]
+    fp16: bool
+
+    def to_dict(self):
+        result = self.__dict__
+        # For serialization, it's best to convert Enums to strings (or their underlying value type).
+        for key, value in result.items():
+            if isinstance(value, Enum):
+                result[key] = value.value
+        return result
+
+    @classmethod
+    def from_json_file(cls, json_file=None):
+        json_file = default_json_config_file if json_file is None else json_file
+        with open(json_file, "r", encoding="utf-8") as f:
+            return cls(**json.load(f))
+
+    def to_json_file(self, json_file):
+        with open(json_file, "w", encoding="utf-8") as f:
+            content = json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+    @classmethod
+    def from_yaml_file(cls, yaml_file=None):
+        yaml_file = default_yaml_config_file if yaml_file is None else yaml_file
+        with open(yaml_file, "r", encoding="utf-8") as f:
+            return cls(**yaml.safe_load(f))
+
+    def to_yaml_file(self, yaml_file):
+        with open(yaml_file, "w", encoding="utf-8") as f:
+            yaml.safe_dump(self.to_dict(), f)
+
+    def __post_init__(self):
+        if isinstance(self.distributed_type, str):
body: Would also put those two functions in the intermediate init `accelerate.commands.config.__init__.py`
diff_hunk: @@ -23,8 +23,8 @@
 from pathlib import Path
 from typing import Optional
 
-from accelerate.commands.config import LaunchConfig, default_config_file
-from accelerate.state import DistributedType
+from accelerate.commands.config.config_args import default_config_file, load_config_from_file
body: I think `LOCAL_MACHINE` is better.
diff_hunk: @@ -53,6 +53,38 @@ class DistributedType(str, Enum):
     TPU = "TPU"
 
 
+class SageMakerDistributedType(str, Enum):
+    """
+    Represents a type of distributed environment.
+
+    Values:
+
+        - **NO** -- Not a distributed environment, just a single process.
+        - **DATA_PARALLEL** -- using sagemaker distributed data parallelism.
+        - **MODEL_PARALLEL** -- using sagemaker distributed model parallelism.
+    """
+
+    # Subclassing str as well as Enum allows the `SageMakerDistributedType` to be JSON-serializable out of the box.
+    NO = "NO"
+    DATA_PARALLEL = "DATA_PARALLEL"
+    MODEL_PARALLEL = "MODEL_PARALLEL"
+
+
+class ComputeEnvironment(str, Enum):
+    """
+    Represents a type of the compute environment.
+
+    Values:
+
+        - **CUSTOM_CLUSTER** -- private/custom cluster hardware.
+        - **AMAZON_SAGEMAKER** -- Amazon SageMaker as compute environment.
+    """
+
+    # Subclassing str as well as Enum allows the `ComputeEnvironment` to be JSON-serializable out of the box.
+    CUSTOM_CLUSTER = "CUSTOM_CLUSTER"
body: Would define a public `is_boto3_available()` function instead of a private variable.
diff_hunk: @@ -25,6 +26,8 @@
 if is_tpu_available():
     import torch_xla.core.xla_model as xm
 
+_has_boto3 = importlib.util.find_spec("boto3") is not None
+
body: I think we can add an extra "sagemaker" with boto3 inside.
diff_hunk: @@ -17,8 +17,17 @@
 
 extras = {}
 extras["quality"] = ["black >= 20.8b1", "isort >= 5.5.4", "flake8 >= 3.8.3"]
-extras["docs"] = ["recommonmark", "sphinx==3.2.1", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
-
+extras["docs"] = [
+    "recommonmark",
+    "sphinx==3.2.1",
+    "sphinx-markdown-tables",
+    "sphinx-rtd-theme==0.4.3",
+    "sphinx-copybutton",
+]
+extras["test"] = [
+    "pytest",
+    "pytest-xdist",
+]
 setup(

Discussion: body: I can´t reproduce why `quality` is failing. I created a complete new conda env with the same packages.