Fix AutoTP custom patterns: respect use_default_specs

[tohtana]

The current code has the following issues:

• use_default_specs: false doesn't work
• Injection by the traditional pattern runs even when custom patterns are set
• mpu needs to be passed to deepspeed.initialize (HF integration doesn't pass mpu)

This PR fixes AutoTP setup to respect use_default_specs: false and disable the traditional injection path when custom patterns are enabled. Also, when mpu is not passed, we create a TP group in the initialization process.

With these changes, the related tests pass and all AutoTP examples in DeepSpeedExamples work now (PR).

[delock]

Are these fix exposed by a test? i.e. a model with conv linear layer or fused qkv weight.

[tohtana]

Great catch! I added a test to validate we use the layers for new custom patterns when a partition is given.

[delock]

Hi @tohtana I reviewed this PR. I have one extended question. I saw current partition_config has a field called "use_default_specs" which allows user to specify default behavior, for example taken from your test:

       partition_config = {
            "use_default_specs":
            False,
            "layer_specs": [
                {
                    "patterns": [".*linears\\.0\\.weight$"],
                    "partition_type": "row",
                },
                {
                    "patterns": [".*linears\\.1\\.weight$"],
                    "partition_type": "column",
                },
            ],
        }

Does this equal to define a default fallback pattern and partition type in layer_specs? I feel this might make the config in a more unified format. It might just be personal favor.

       partition_config = {
            #"use_default_specs":
            #False,
            "layer_specs": [
                {
                    "patterns": [".*linears\\.0\\.weight$"],
                    "partition_type": "row",
                },
                {
                    "patterns": [".*linears\\.1\\.weight$"],
                    "partition_type": "column",
                },
                { # the default fallback pattern and partition type
                    "patterns": ["*"],
                    "partition_type": "none",
                },
            ],
        }

[tohtana]

@delock Thank you for your review!

I saw current partition_config has a field called "use_default_specs" which allows user to specify default behavior, for example taken from your test:

use_default_specs doesn't fallback to "partition_type": "none". It controls whether we use preset or not.
With use_default_specs=True, we will add the custom patterns to the preset patterns.
When it's False, only the custom patterns in the config are used (This assumes we disable the preset when it is enabled by a different stack)

The document explains "You can also set use_default_specs to true to merge your custom patterns on top of the preset (when preset_model is provided)." I think we could clarify this by explaining when use_default_specs=False is useful.
We can also add a warning when preset_model and use_default_specs=False are both given.

Can you share your thoughts?

[delock]

@tohtana Thanks for the explaination! I don't have further comments.

@delock Thank you for your review!

I saw current partition_config has a field called "use_default_specs" which allows user to specify default behavior, for example taken from your test:

use_default_specs doesn't fallback to "partition_type": "none". It controls whether we use preset or not. With use_default_specs=True, we will add the custom patterns to the preset patterns. When it's False, only the custom patterns in the config are used (This assumes we disable the preset when it is enabled by a different stack)

The document explains "You can also set use_default_specs to true to merge your custom patterns on top of the preset (when preset_model is provided)." I think we could clarify this by explaining when use_default_specs=False is useful. We can also add a warning when preset_model and use_default_specs=False are both given.

Can you share your thoughts?