Stop the curriculum schedule starting below min_difficulty (#8334)

## Symptom

The curriculum schedule can start below the `min_difficulty` it is
configured with, and for one pair the tutorial itself recommends it
starts at **0**.

```python
from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler

def sched(min_d, step):
    return CurriculumScheduler({
        "min_difficulty": min_d, "max_difficulty": 1024, "schedule_type": "fixed_linear",
        "schedule_config": {"total_curriculum_step": 100, "difficulty_step": step},
    })

for min_d, step in [(8, 8), (8, 16), (64, 16), (1, 8), (10, 8)]:
    s = sched(min_d, step)
    print(min_d, step, [s.get_difficulty(i) for i in range(6)])
```

| `min_difficulty` | `difficulty_step` | first six steps |
| --- | --- | --- |
| 8 | 8 | `[8, 16, 24, 32, 48, 56]` |
| **8** | **16** | **`[0, 16, 16, 32, 48, 48]`** |
| 64 | 16 | `[64, 64, 80, 80, 96, 112]` |
| **1** | **8** | **`[0, 8, 16, 24, 40, 48]`** |
| **10** | **8** | **`[8, 16, 24, 40, 48, 56]`** |

`min_difficulty=8` with `difficulty_step=16` is not a contrived pair.
The tutorial recommends "starting with `min_difficulty` at 8
(million-scale models) or 64 (billion-scale models)" and separately "we
usually set [`difficulty_step`] to 8 (for FP16 data) or 16 (for INT8
data)". A million-scale model on INT8 data lands on exactly that
combination, and its first training step gets a sequence length of 0.

## Root cause

`__fixed_root_get_difficulty`, which serves both `fixed_linear` (root
degree 1) and `fixed_root`, floors the interpolated value to a multiple
of `difficulty_step` and then clamps only the top:

```python
next_difficulty -= (next_difficulty % s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP])
next_difficulty = min(next_difficulty, self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY])
```

At step 0 the interpolation is exactly `min_difficulty`, so the floor
subtracts `min_difficulty % difficulty_step` and there is nothing to
stop it going under. The tutorial's own formula for this schedule is
`((step/total)**(1/root_degree)) * (max_difficulty - min_difficulty) +
min_difficulty`, which starts at `min_difficulty`.

## Fix

Clamp the bottom the way the top already is, one line.

This does not introduce a new exception to the "difficulty is a multiple
of `difficulty_step`" rule: the existing top clamp already returns
`max_difficulty` verbatim when it is not a multiple. With
`max_difficulty=1000` and `difficulty_step=16` the schedule returns
1000, not 992, once it runs out. Both endpoints being the configured
values rather than multiples of the step is the behaviour this function
already has at one end.

`__fixed_discrete_get_difficulty` picks from an explicit list and is
untouched.

## Test

Two tests in `tests/unit/runtime/test_data_efficiency.py`, both plain
CPU tests rather than `DistributedTest`, since `CurriculumScheduler`
needs neither an accelerator nor a process group:

- `test_curriculum_never_starts_below_min_difficulty`, parametrized over
`fixed_linear` and `fixed_root` and over five `(min_difficulty,
difficulty_step)` pairs including the tutorial's own recommendations,
checks the first twenty steps stay within `[min_difficulty,
max_difficulty]`.
- `test_curriculum_endpoints_are_the_configured_values` pins both ends
with a `max_difficulty` that is not a multiple of `difficulty_step`. Its
top-end assertion passes on master too, which is what makes it the
control for the argument above.

Against master: **7 failed, 4 passed, 6 skipped** (`assert 0 >= 8`,
`assert 8 >= 10`, `assert 0 == 8`). With the fix: **11 passed, 6
skipped**. The 6 skipped are the file's pre-existing `DistributedTest`
cases, which need 2 GPUs; a pristine checkout reports the same 6 skips
and nothing else.

yapf and flake8 clean, with yapf making no changes to either file.

---------

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

Author: vineethsaivs

deepspeed/runtime/data_pipeline/curriculum_scheduler.py

@@ -137,6 +137,15 @@ def __fixed_root_get_difficulty(self, global_steps, root_degree=None):
             (self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY] - self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY]) +
             self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY])
         next_difficulty -= (next_difficulty % s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP])
+        # Flooring to a multiple of difficulty_step can land below the configured start, and
+        # does for pairs the tutorial itself recommends: min_difficulty 8 with difficulty_step
+        # 16 gives 0 on the first step, which is a zero-length sequence for the seqlen metric.
+        # Raise the floor to the first multiple of the step at or above min_difficulty rather
+        # than to min_difficulty itself: the constructor warns that every difficulty has to be
+        # a multiple of the step (8 for FP16, 16 for INT8), and a min_difficulty that is not
+        # one cannot also be the first difficulty. This way both contracts hold.
+        step = s_state[CURRICULUM_LEARNING_SCHEDULE_DIFFICULTY_STEP]
+        next_difficulty = max(next_difficulty, -(-self.state[CURRICULUM_LEARNING_MIN_DIFFICULTY] // step) * step)
         next_difficulty = min(next_difficulty, self.state[CURRICULUM_LEARNING_MAX_DIFFICULTY])
         return next_difficulty
 

tests/unit/runtime/test_data_efficiency.py

@@ -10,6 +10,7 @@
 import pytest
 from unit.common import DistributedTest
 from unit.simple_model import Curriculum_SimpleModel, SimpleModel, random_dataloader, random_dataset
+from deepspeed.runtime.data_pipeline.curriculum_scheduler import CurriculumScheduler
 
 
 class MPU():
@@ -50,6 +51,54 @@ def get_model_parallel_group(self):
         return self.tp_group
 
 
+def _curriculum_scheduler(min_difficulty, max_difficulty, difficulty_step, schedule_type="fixed_linear"):
+    config = {
+        "min_difficulty": min_difficulty,
+        "max_difficulty": max_difficulty,
+        "schedule_type": schedule_type,
+        "schedule_config": {
+            "total_curriculum_step": 100,
+            "difficulty_step": difficulty_step,
+        },
+    }
+    if schedule_type == "fixed_root":
+        config["schedule_config"]["root_degree"] = 2
+    return CurriculumScheduler(config)
+
+
+@pytest.mark.parametrize("schedule_type", ["fixed_linear", "fixed_root"])
+@pytest.mark.parametrize("min_difficulty, difficulty_step", [(8, 16), (1, 8), (10, 8), (64, 16), (8, 8), (100, 64)])
+def test_curriculum_never_starts_below_min_difficulty(schedule_type, min_difficulty, difficulty_step):
+    # Rounding down to a multiple of difficulty_step used to push the first steps under
+    # the configured start: min_difficulty 8 with difficulty_step 16, a pair the tutorial
+    # recommends for INT8 data on a million-scale model, gave a difficulty of 0, which is
+    # a zero-length sequence for the seqlen metric. The step alignment the constructor
+    # warns about has to survive the new floor, so it is asserted alongside.
+    scheduler = _curriculum_scheduler(min_difficulty, 1024, difficulty_step, schedule_type)
+    difficulties = [scheduler.get_difficulty(step) for step in range(100)]
+    assert min(difficulties) >= min_difficulty
+    assert max(difficulties) <= 1024
+    assert all(difficulty % difficulty_step == 0 for difficulty in difficulties)
+
+
+def test_curriculum_first_difficulty_is_the_aligned_min():
+    # min_difficulty 8 is not a multiple of difficulty_step 16, so it cannot itself be a
+    # difficulty. The schedule starts at the first multiple at or above it, 16, which is
+    # the only value that is both no lower than asked and aligned to the step.
+    scheduler = _curriculum_scheduler(8, 1024, 16)
+    assert scheduler.get_difficulty(0) == 16
+
+    # a min_difficulty that is already a multiple is used unchanged
+    assert _curriculum_scheduler(64, 1024, 16).get_difficulty(0) == 64
+
+
+def test_curriculum_tops_out_at_the_configured_max():
+    # max_difficulty is not a multiple of difficulty_step here. The existing clamp caps
+    # the ramp at the configured value, and this change does not touch that end.
+    scheduler = _curriculum_scheduler(16, 1000, 16)
+    assert scheduler.get_difficulty(200) == 1000
+
+
 @pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float16])
 class TestDataEfficiency(DistributedTest):
     world_size = 2