Keep the elasticity batch overrides out of the caller's config dict

[alanhuangyoo]

DeepSpeedConfig keeps the dict it is handed by reference:

if isinstance(config, dict):
    self._param_dict = config

#8289 established that parsing must not write back into it — the caller owns that dict and may reuse it after initialization.

The elasticity branch still does, two lines above the comment that says otherwise:

        self._param_dict[TRAIN_BATCH_SIZE] = final_batch_size
        self._param_dict[TRAIN_MICRO_BATCH_SIZE_PER_GPU] = micro_batch_size
        self._param_dict[GRADIENT_ACCUMULATION_STEPS] = gradient_accu_steps

    # Pass a copy so that user json is unmodified, e.g. for logging
    self._initialize_params(copy.copy(self._param_dict))

A caller that enables elasticity gets three keys back that it never set. print_user_config() dumps self._param_dict, so it then reports them as though the user had written them.

It also makes the dict unparseable a second time. The elasticity path rejects those keys in the input unless ignore_non_elastic_batch_info is set:

One or more batch related parameters were found in your ds_config (...).
These parameters *will not be used* since elastic training is enabled ...

The first parse succeeds and injects them; the second parse of the same dict trips that guard, and its message asks the user to remove three keys they never wrote.

The fix

Collect the overrides and apply them to the copy. All three are top-level keys, so the existing shallow copy keeps them off the caller's dict.

Test

DeepSpeedConfig(config_dict) twice on an elasticity config, with ignore_non_elastic_batch_info left out so the guard is live:

before

parse 1 OK, caller dict gained: ['gradient_accumulation_steps', 'train_batch_size', 'train_micro_batch_size_per_gpu']
parse 2 FAILED: ElasticityConfigError One or more batch related parameters were found in your ds_config ...

after

parse 1 OK, caller dict gained: nothing
parse 2 OK

Parsed values are unchanged either way (train_batch_size=4 micro=2 gas=2), so this only removes the write-back.

test_elasticity_leaves_caller_config_untouched sits next to #8289's test_max_grad_norm_leaves_caller_config_untouched and covers both symptoms. On master it fails at

AssertionError: assert {'elasticity', 'train_batch_size', 'train_micro_batch_size_per_gpu',
                        'gradient_accumulation_steps'} == {'elasticity'}

tests/unit/runtime/test_ds_config_dict.py   27 passed, 5 skipped
tests/unit/elasticity/test_elastic.py       23 passed, 3 skipped
yapf --diff / flake8                        clean

[ebarkhordar]

I ran the double-parse case against b0dee8d and its merge base 32e301f in a clean container. The change looks right to me. The re-parse sentence in the description has the direction backwards.

It says a config "that was rejected on the first pass would be accepted on a re-parse". The guard at config.py:764 raises when the batch keys are present, so it is the other way round: the first parse succeeds and injects them, and the second parse of the same dict is the one that fails.

Elasticity on, ignore_non_elastic_batch_info unset, DeepSpeedConfig(d) twice on the same dict:

32e301ffa  parse 1 OK, caller dict gains the 3 batch keys
           parse 2 ElasticityConfigError: One or more batch related parameters were found in your ds_config (...)
b0dee8df2  parse 1 OK, caller dict unchanged
           parse 2 OK

That is a stronger case for the fix than the description makes: on master, re-parsing a dict you just parsed is a hard error, and its text asks the user to remove three keys they never wrote.

It is also unpinned. test_elasticity_leaves_caller_config_untouched sets ignore_non_elastic_batch_info: True, which switches that guard off, so it cannot reach the path. Leaving the flag out and calling DeepSpeedConfig(config_dict) twice covers it.

I exercised only the dict branch, not the json path.

[alanhuangyoo]

You are right on both, thanks — and the repro is better than what I wrote.

I had the direction backwards. The guard at config.py:764 fires when the batch keys are present, so the first parse is the one that succeeds and injects them, and the second parse of the same dict is the one that dies. I reproduced yours:

master     parse 1 OK, caller dict gained: [gradient_accumulation_steps, train_batch_size, train_micro_batch_size_per_gpu]
           parse 2 FAILED: ElasticityConfigError One or more batch related parameters were found in your ds_config ...
this PR    parse 1 OK, caller dict gained: nothing
           parse 2 OK

Description rewritten around that: on master, re-parsing a dict you just parsed is a hard error whose message names three keys the caller never wrote.

The test was unpinned exactly as you say — ignore_non_elastic_batch_info: True skips the guard, so the second parse could not fail and the test would have passed with or without the write-back. Dropped the flag and it now parses the same dict twice. On master it fails at the key-set assertion; if someone were to fix only that symptom and still write into _param_dict, the second parse catches it.

Left the json branch alone — same as you, I only exercised the dict path.

[ebarkhordar]

Confirmed at b6cc16f8. Dropping the flag makes the test load-bearing: it passes at the PR head and fails at the merge base 32e301ff with head's test file copied in.

pytest tests/unit/runtime/test_ds_config_dict.py::test_elasticity_leaves_caller_config_untouched

head b6cc16f8               1 passed
base 32e301ff + this test   1 failed
  test_ds_config_dict.py:293: AssertionError
  assert {'elasticity', 'gradient_accumulation_steps', 'train_batch_size',
          'train_micro_batch_size_per_gpu'} == {'elasticity'}

It stops at line 293, so the second parse is unreached at base, which matches your reading of it as the backstop rather than the primary assertion. I only ran the dict path, same as you.

[alanhuangyoo]

The red on modal-torch-latest is not from this change:

FAILED tests/unit/v1/compile/test_offload_activation.py::test_floor_peak_agrees_across_ranks

That lane passed on b0dee8d (44m56s), and the only thing b6cc16f8 adds is 13 lines in tests/unit/runtime/test_ds_config_dict.py — test_ds_config_dict does not appear anywhere in the failing run. A cross-rank peak-memory agreement check in activation offloading has no path to a config-parsing test.

The lane is green on master, so this looks like a flake in that test rather than a broken lane. Happy to rebase if a re-run is easier than taking my word for it.

The rest of the lanes on this PR are sitting in action_required.