add distributed config - #46705
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
View the CircleCI Test Summary for this PR: https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=46705&sha=799ac9 |
|
CI Dashboard: View test results in Grafana |
ArthurZucker
left a comment
There was a problem hiding this comment.
would remove the test here imo
| # easily available | ||
| self._tp_plan, self._ep_plan, self._pp_plan = {}, {}, {} | ||
| # easily available. | ||
| self._tp_plan, self._ep_plan, self._pp_plan, self._fsdp_plan = {}, {}, {}, {} |
There was a problem hiding this comment.
don't you want to add the mixin in this PR?
There was a problem hiding this comment.
I prefer to delay it to another PR once more thing have been merged. This way, it will be easier to know what will be abstracted away in the DistributedMixin once everything is there
| @@ -0,0 +1,97 @@ | |||
| import json | |||
| class TestDistributedConfig: | ||
| def test_2d_parallelism(self): | ||
| dc = DistributedConfig(tp_size=2, fsdp_size=2) | ||
| assert dc.tp_size == 2 | ||
| assert dc.fsdp_size == 2 | ||
| assert dc.tp_plan is None | ||
| assert dc.fsdp_cpu_offload is False | ||
| assert dc.fsdp_mixed_precision is False | ||
|
|
||
| def test_tp_only_defaults_fsdp_to_1(self): | ||
| dc = DistributedConfig(tp_size=4) | ||
| assert dc.tp_size == 4 | ||
| assert dc.fsdp_size == 1 | ||
| assert dc.tp_plan is None | ||
|
|
||
| def test_fsdp_only_defaults_tp_to_1(self): | ||
| dc = DistributedConfig(fsdp_size=4) | ||
| assert dc.tp_size == 1 | ||
| assert dc.fsdp_size == 4 | ||
| assert dc.fsdp_cpu_offload is False | ||
| assert dc.fsdp_mixed_precision is False | ||
| assert dc.tp_plan is None | ||
|
|
||
| def test_empty_config(self): | ||
| dc = DistributedConfig() | ||
| assert dc.tp_size is None | ||
| assert dc.fsdp_size is None | ||
| assert dc.tp_plan is None | ||
| assert dc.fsdp_cpu_offload is False | ||
| assert dc.fsdp_mixed_precision is False | ||
|
|
||
| def test_from_dict(self): | ||
| dc = DistributedConfig.from_dict({"tp_size": 2, "fsdp_size": 4}) | ||
| assert dc.tp_size == 2 | ||
| assert dc.fsdp_size == 4 | ||
| assert dc.tp_plan is None | ||
|
|
||
| def test_from_dict_ignores_unknown_keys(self): | ||
| dc = DistributedConfig.from_dict({"tp_size": 2, "unknown_key": 42}) | ||
| assert dc.tp_size == 2 | ||
| assert not hasattr(dc, "unknown_key") | ||
|
|
||
| def test_from_dict_kwargs_override(self): | ||
| dc = DistributedConfig.from_dict({"tp_size": 2}, tp_size=8) | ||
| assert dc.tp_size == 8 | ||
|
|
||
| def test_to_dict(self): | ||
| dc = DistributedConfig(tp_size=2, fsdp_size=4) | ||
| d = dc.to_dict() | ||
| assert d == { | ||
| "tp_size": 2, | ||
| "tp_plan": None, | ||
| "enable_sequence_parallel": False, | ||
| "enable_expert_parallel": False, | ||
| "fsdp_size": 4, | ||
| "fsdp_cpu_offload": False, | ||
| "fsdp_mixed_precision": False, | ||
| } | ||
|
|
||
| def test_to_dict_is_a_copy(self): | ||
| dc = DistributedConfig(tp_plan={"layer": "colwise"}) | ||
| d = dc.to_dict() | ||
| d["tp_plan"]["layer"] = "rowwise" | ||
| assert dc.tp_plan["layer"] == "colwise" | ||
|
|
||
| def test_to_json_string(self): | ||
| dc = DistributedConfig(tp_size=2, fsdp_size=2) | ||
| s = dc.to_json_string() | ||
| parsed = json.loads(s) | ||
| assert parsed["tp_size"] == 2 | ||
| assert parsed["fsdp_size"] == 2 | ||
|
|
||
| def test_to_json_file(self): | ||
| dc = DistributedConfig(tp_size=4) | ||
| with tempfile.NamedTemporaryFile(mode="r", suffix=".json", delete=False) as f: | ||
| dc.to_json_file(f.name) | ||
| f.seek(0) | ||
| parsed = json.load(f) |
There was a problem hiding this comment.
don't think we need such granularity!
* add distributed config * fix * fix * remove redundant test file
No description provided.