2023 the huggingface team 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 orglicenseslicense2 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 default torch distributed port when using a single gpu launcher emulation i e not deepspeed or python m torch distributed the issue is that once the port is tied it can t be used anywhere else outside of this process since torch dist doesn t free the port until the process exits therefore for the sake of being able to run both emulated launcher and normal launcher tests we need 2 distinct ports this function will give the right port in the right context for real launcher it ll give the base port for emulated launcher it ll give the base port 1 in both cases a string is returned args reallauncher whether a real launcher is going to be used or the emulated one hack to restore original logging level pre 21700 1 explicitly set numnodes1 just in case these tests end up run on a multinode setup it won t be able to handle that 2 for now testing with just 2 gpus max since some quality tests may give different results with mode gpus because we use very little data customize the test name generator function as we want both params to appear in the subtest name as by default it shows only the first param resume from ckpt split cmd launcher fsdpconfig script args keep for quick debug print joinfnpythonpathself srcdirstr cmd die executesubprocessasynccmd envself getenv logs trainerstate loadfromjsonos path joinoutputdir trainerstate json loghistory return logs def getbaseargsself outputdir numepochs loggingsteps return f modelnameorpath bertbasecased taskname mrpc outputdir outputdir overwriteoutputdir dotrain maxseqlength 128 perdevicetrainbatchsize 16 learningrate 5e5 numtrainepochs numepochs lrschedulertype cosine loggingsteps loggingsteps savestrategy epoch doeval evaluationstrategy epoch reportto none 2023 the huggingface team 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 noqa default torch distributed port when using a single gpu launcher emulation i e not deepspeed or python m torch distributed the issue is that once the port is tied it can t be used anywhere else outside of this process since torch dist doesn t free the port until the process exits therefore for the sake of being able to run both emulated launcher and normal launcher tests we need 2 distinct ports this function will give the right port in the right context for real launcher it ll give the base port for emulated launcher it ll give the base port 1 in both cases a string is returned args real_launcher whether a real launcher is going to be used or the emulated one noqa hack to restore original logging level pre 21700 1 explicitly set num_nodes 1 just in case these tests end up run on a multi node setup it won t be able to handle that 2 for now testing with just 2 gpus max since some quality tests may give different results with mode gpus because we use very little data accelerate launch num_processes num_gpus main_process_port master_port use_fsdp fsdp_auto_wrap_policy transformer_based_wrap fsdp_state_dict_type sharded_state_dict fsdp_transformer_layer_cls_to_wrap bertlayer customize the test name generator function as we want both params to appear in the sub test name as by default it shows only the first param resume from ckpt fsdp_sharding_strategy fsdp_sharding_strategy index sharding_strategy upper 1 keep for quick debug print join f npythonpath self src_dir_str cmd die model_name_or_path bert base cased task_name mrpc output_dir output_dir overwrite_output_dir do_train max_seq_length 128 per_device_train_batch_size 16 learning_rate 5e 5 num_train_epochs num_epochs lr_scheduler_type cosine logging_steps logging_steps save_strategy epoch do_eval evaluation_strategy epoch report_to none
import itertools
import os
import unittest
from functools import partial

from parameterized import parameterized

import tests.trainer.test_trainer
from tests.trainer.test_trainer import TrainerIntegrationCommon  
from transformers import is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    backend_device_count,
    execute_subprocess_async,
    mockenv_context,
    require_accelerate,
    require_fsdp,
    require_torch_accelerator,
    require_torch_multi_accelerator,
    slow,
    torch_device,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import FSDPOption, set_seed
from transformers.utils import is_accelerate_available, is_torch_bf16_available_on_device


if is_torch_available():
    from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_1
else:
    is_torch_greater_or_equal_than_2_1 = False


DEFAULT_MASTER_PORT = "10999"
dtypes = ["fp16"]
if is_torch_bf16_available_on_device(torch_device):
    dtypes += ["bf16"]
sharding_strategies = ["full_shard", "shard_grad_op"]
state_dict_types = ["FULL_STATE_DICT", "SHARDED_STATE_DICT"]
set_seed(42)
params = list(itertools.product(sharding_strategies, dtypes))


def get_master_port(real_launcher=False):
    

    master_port_base = os.environ.get("DS_TEST_PORT", DEFAULT_MASTER_PORT)
    if not real_launcher:
        master_port_base = str(int(master_port_base) + 1)
    return master_port_base


if is_torch_available():
    from tests.trainer.test_trainer import (  
        RegressionModelConfig,
        RegressionPreTrainedModel,
    )

    
    get_regression_trainer = partial(tests.trainer.test_trainer.get_regression_trainer, log_level="info")

require_fsdp_version = require_fsdp
if is_accelerate_available():
    from accelerate.utils.constants import (
        FSDP_PYTORCH_VERSION,
        FSDP_SHARDING_STRATEGY,
    )

    require_fsdp_version = partial(require_fsdp, min_version=FSDP_PYTORCH_VERSION)


def get_launcher(distributed=False, use_accelerate=False):
    
    
    
    
    num_gpus = min(2, backend_device_count(torch_device)) if distributed else 1
    master_port = get_master_port(real_launcher=True)
    if use_accelerate:
        return f.split()
    return f"torchrun --nnodes 1 --nproc-per-node {num_gpus} --master-port {master_port}".split()


def _parameterized_custom_name_func(func, param_num, param):
    
    
    param_based_name = parameterized.to_safe_name("_".join(str(x) for x in param.args))
    return f"{func.__name__}_{param_based_name}"


@require_accelerate
@require_torch_accelerator
@require_fsdp_version
class TrainerIntegrationFSDP(TestCasePlus, TrainerIntegrationCommon):
    def setUp(self):
        super().setUp()
        master_port = get_master_port(real_launcher=False)
        self.dist_env_1_gpu = {
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": master_port,
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
        }

        self.fsdp_config = {
            "backward_prefetch": "backward_pre",
            "forward_prefetch": "False",
            "limit_all_gathers": "False",
            "use_orig_params": "True",
            "sync_module_states": "True",
            "activation_checkpointing": "False",
            "min_num_params": 1,
        }

    def tearDown(self):
        super().tearDown()

    @parameterized.expand(params, name_func=_parameterized_custom_name_func)
    def test_fsdp_config(self, sharding_strategy, dtype):
        output_dir = self.get_auto_remove_tmp_dir()
        kwargs = {
            "output_dir": output_dir,
            "train_len": 128,
            "save_steps": 5,
            "learning_rate": 0.1,
            "fsdp": f"{sharding_strategy} offload auto_wrap",
            "fsdp_config": self.fsdp_config,
        }
        kwargs[dtype] = True
        with mockenv_context(**self.dist_env_1_gpu):
            trainer = get_regression_trainer(**kwargs)
            self.assertEqual(trainer.args.fsdp[0], sharding_strategy)
            self.assertEqual(trainer.args.fsdp[1], FSDPOption.OFFLOAD)
            self.assertEqual(trainer.args.fsdp[2], FSDPOption.AUTO_WRAP)
            for k, v in trainer.args.fsdp_config.items():
                self.assertEqual(v, self.fsdp_config[k])
            self.assertEqual(os.environ.get("ACCELERATE_USE_FSDP", "false"), "true")

    @parameterized.expand(params, name_func=_parameterized_custom_name_func)
    @require_torch_multi_accelerator
    @slow
    def test_basic_run(self, sharding_strategy, dtype):
        launcher = get_launcher(distributed=True, use_accelerate=False)
        output_dir = self.get_auto_remove_tmp_dir()
        args = self.get_base_args(output_dir, 1, 50).split() + [f"--{dtype}"]
        fsdp_args = ["--fsdp", f"{sharding_strategy} auto_wrap", "--fsdp_transformer_layer_cls_to_wrap", "BertLayer"]
        script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
        cmd = launcher + script + args + fsdp_args
        execute_subprocess_async(cmd, env=self.get_env())

    @parameterized.expand(dtypes)
    @require_torch_multi_accelerator
    @slow
    @unittest.skipIf(not is_torch_greater_or_equal_than_2_1, reason="This test on pytorch 2.0 takes 4 hours.")
    def test_basic_run_with_cpu_offload(self, dtype):
        launcher = get_launcher(distributed=True, use_accelerate=False)
        output_dir = self.get_auto_remove_tmp_dir()
        args = self.get_base_args(output_dir, 1, 50).split() + [f"--{dtype}", "--max_steps", "10"]
        fsdp_args = ["--fsdp", "full_shard auto_wrap offload", "--fsdp_transformer_layer_cls_to_wrap", "BertLayer"]
        script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
        cmd = launcher + script + args + fsdp_args
        execute_subprocess_async(cmd, env=self.get_env())

    @parameterized.expand(state_dict_types, name_func=_parameterized_custom_name_func)
    @require_torch_multi_accelerator
    @slow
    def test_training_and_can_resume_normally(self, state_dict_type):
        output_dir = self.get_auto_remove_tmp_dir("./xxx", after=False)

        sharding_strategy = "full_shard"
        use_accelerate = state_dict_type == "SHARDED_STATE_DICT"
        launcher = get_launcher(True, use_accelerate=use_accelerate)
        args = self.get_base_args(output_dir, 2, 25).split()
        script = [f"{self.examples_dir_str}/pytorch/text-classification/run_glue.py"]
        logs = self.run_cmd_and_get_logs(use_accelerate, sharding_strategy, launcher, script, args, output_dir)

        
        checkpoint = os.path.join(output_dir, "checkpoint-115")
        resume_args = args + f"--resume_from_checkpoint {checkpoint}".split()
        logs_resume = self.run_cmd_and_get_logs(
            use_accelerate, sharding_strategy, launcher, script, resume_args, output_dir
        )

        for log, log1 in zip(logs, logs_resume):
            if "learning_rate" in log:
                self.assertAlmostEqual(log["learning_rate"], log1["learning_rate"], delta=1e-5)

    def run_cmd_and_get_logs(self, use_accelerate, sharding_strategy, launcher, script, args, output_dir):
        if not use_accelerate:
            fsdp_args = [
                "--fsdp",
                f"{sharding_strategy} auto_wrap",
                "--fsdp_transformer_layer_cls_to_wrap",
                "BertLayer",
            ]
            cmd = launcher + script + args + fsdp_args
        else:
            fsdp_config = f.split()
            cmd = launcher + fsdp_config + script + args

        
        
        execute_subprocess_async(cmd, env=self.get_env())
        logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history
        return logs

    def get_base_args(self, output_dir, num_epochs, logging_steps):
        return f