possibly --skip-train-iteration-range with multiple entries has a bug

[stas00]

Just noticed in the logs that --skip-train-iteration-range reports only a single range, when there should be 2:

I currently have in the config:

--skip-train-iteration-range 13251-14000 16651-19500

But the report is only:

Skipped iterations 16651 to 19500 due to --skip-train-iteration-range flag.

@jaketae, could you please check if it's just printing error or whether it actually doesn't work as we planned it to? Thanks!

[jaketae]

That's odd, thanks for flagging! I'll look over it over the weekend.

[jaketae]

@stas00 Do we know what the iteration was when you ran the script with the skip flags? My best educated guess is that this is the culprit:

Megatron-DeepSpeed/megatron/training.py

Lines 799 to 804 in 04c461e

	# flush intervals prior to current iteration
	if args.skip_train_iteration_range is not None:
	ends = [end for start, end in args.skip_train_iteration_range]
	index = bisect.bisect_left(ends, iteration)
	for _ in range(index):
	args.skip_train_iteration_range.popleft()

Is it possible that the the iteration count was larger than the first skip range when training was resumed?

[stas00]

yes, it will always be larger as the training progresses, yet we need to continue skipping earlier ranges.

so with --skip-train-iteration-range 13251-14000 16651-19500 the actual iteration will typically be somewhere past 16651.

But also say we want to re-train the whole thing keeping these ranges, then iteration can be anything - smaller and larger.

Does it make sense?

So I think if the current iteration is larger than the first range, it should always cycle through the first range.

And if it's already in the second range, then effectively it should cycle through the sum of the two.

It's a tricky one.

When we have --skip-train-iteration-range 13251-14000 16651-19500, the first range is well defined, but the second range really means 16651+(14000-13251)-19500+(14000-13251) and so on if there are more ranges.

So perhaps if following the summing logic from the para above then it actually wouldn't be much different from your current code, just add up a sum of any previous ranges to the current range that the iteration is overlapping or just passed it.

[jaketae]

But also say we want to re-train the whole thing keeping these ranges, then iteration can be anything - smaller and larger.

This totally makes sense. But if the skip range was 13251-14000 16651-19500 and the loaded checkpoint started at, say iteration 16000, the first skip range is meaningless because it's already in the past and can't be skipped, which is what I suspect might have happened. Am I missing something?

So I think if the current iteration is larger than the first range, it should always cycle through the first range.

I'm not sure if I follow... Could you perhaps elaborate a bit more please?

When we have --skip-train-iteration-range 13251-14000 16651-19500, the first range is well defined, but the second range really means 16651+(14000-13251)-19500+(14000-13251) and so on if there are more ranges.

This makes sense, and I think we accounted for this by incrementing a dummy variable iteration_for_skipping instead of the actual iteration, which is used to determine whether to skip a micro batch.

Megatron-DeepSpeed/megatron/training.py

Lines 811 to 822 in 04c461e

	and args.skip_train_iteration_range[0][0] <= iteration + 1 <= args.skip_train_iteration_range[0][1]
	):
	start, end = args.skip_train_iteration_range.popleft()
	print_rank_0(f"Skipped iterations {start} to {end} due to --skip-train-iteration-range flag.")
	iteration_for_skipping = args.iteration
	while iteration_for_skipping + 1 <= end:
	try:
	_ = next(train_data_iterator)
	except TypeError:
	pass
	iteration_for_skipping += 1
	continue

I feel like we had a similar discussion way before, I'll look through the merged PR to see if I can recall some of the bits of thought. Sorry if I'm making you repeat!

[stas00]

But also say we want to re-train the whole thing keeping these ranges, then iteration can be anything - smaller and larger.

This totally makes sense. But if the skip range was 13251-14000 16651-19500 and the loaded checkpoint started at, say iteration 16000, the first skip range is meaningless because it's already in the past and can't be skipped, which is what I suspect might have happened. Am I missing something?

It depends on how you look at it.

Consider this dataset sequence: 123456789

And your skip range is: 2-3 5-6 8-9

Current code with 2 iteration per training cycle will pick:

iteration: 12, 34, 56
data:      14, 47, 71

As you can see on job #2, it'll skip 2-3 and start with 4, but it has already run sample 4.

but we want 14,71,14,71 - no repetition of the same samples until the epoch has run out.

So if you take into consideration a sub-range but not always then you end up with your data being inconsistent.

Each range should basically remove those pockets that should be excluded. So in the example above:

The only valid samples are 1,4,7 and then it's easy to not get overlaps.

Now your look says, look we don't care what happened in the past, we want to skip this particular range that we don't want right now and once we skipped it let's forget it ever was skipped and look at the future data as nothing was skipped.

Both looks are valid and my contrived example becomes a problem only when ranges lead to an overlap of skipped ranges on a new job.

From the user point of view I'm still trying to wrap my head around what --skip-train-iteration-range 13251-14000 16651-19500 does really mean. I thought we were skipping these 2 ranges: 13251-14000 17400-20249 but it proved to be we only skip the last range.

Which on the other hand could be said is more intuitive as long as the overlap condition doesn't occur, since we want to skip data at the currently relevant range and not really care of all the previous windows that we skipped.

So perhaps we can leave it as is for now, but I'd only request to explicitly print out that the other ranges were ignored (and to list them), like:

Skipped iterations 16651 to 19500 due to --skip-train-iteration-range flag. Additionally, iterations 13251-14000 
were skipped in the previous runs.

How does that sound? Here we let the user know that the other numbers haven't been missed, but they are ignored by design.

I think one more issue is that we consider that the dataset wraps around. However if it ends with StopIteration when it runs out of data, this would be very confusing because the current skip range won't be matching the actual dataset length.

---

Also have we considered the scenario of number of iterations per run being shorter than skip range?

Say, you have 30 samples and your skip range is 11-20 and each run is 5 iterations long. So we have:

run | iter   | samples
----|--------|------
  1 | 1-5    | 1-5
  2 | 6-10   | 6-10
  3 | 10-15  | 21-25
  4 | 16-20  | ????

obviously we want 26-30 to be the last batch of samples, but won't we get 21-25 again here a second time as it'd skip the range and start right after it was skipped?

I haven't checked what the code will do - this is just an case scenario that I have thought of while contemplating this problem.

[jaketae]

Sorry @stas00, somehow this slipped under my radar. Has this issue been resolved or fallen out of relevance? If not, I can look at it this week and hopefully ship a fix. Thank you!

[stas00]

I think it may be pretty essential should we run into instabilities during 200B training. It hasn't been resolved or fallen out of relevance. It's waiting for your magic touch, Jaesung :)

[jaketae]

I'm no magician, but you sure are a wizard!

I can certainly add a print that informs the user of any "flushing" of past intervals that has occurred. But before I do that, I want to make sure that I've wrapped my head around everything and that we are on the same page.

1. The script has no way of knowing what happened before. That is, the script doesn't know whether the current run is the first time skipping, the second time, etc. Thus, the burden is on the user to recalculate and update skip intervals on subsequent runs. This probably involves subtracting future skip intervals with the total number of skips that have already been performed on previous runs.
2. In this sense, skip intervals are not absolute indices. They are numbers relative to the current iteration, which itself maybe have been affected by skips on previous runs.
3. If a skip interval is in the relative past, we can ignore it. If the current index is x and the skip interval is from x-n to x-m, this interval is irrelevant since it is in the relative past.
4. Nevertheless, we inform the user of irrelevant past intervals for informative logging. The motivation behind this PR was that it is impossible to know whether a skip was purposefully ignored by the algorithm, or whether it is an unaccounted edge case.

Would you say this is a fair characterization of the problem at hand? Thank you!

[stas00]

I didn't validate every clause you mentioned but your diagnostics sounds correct, Jaesung.

Perhaps the issue stems from the fact that:

(1) we decided to skip iterations instead of samples because it's somewhat easier to use in the moment of wanting to skip (albeit we know the samples count)
(2) that we are not doing the long-term accounting and thus have "memory" gaps.

What if instead of skipping iterations we were to skip samples. Then they can be absolute and we will would always know exactly what has been skipped.

This would require 2 counters - samples_consumed (the existing one) and samples_pulled (a new one). So samples_pulled will be the absolute pointer into the dataset, whereas samples_consumed will continue counting the total samples the model was trained on.

So --skip-samples-range will now interact only with samples_pulled and when a user wants to add a range to skip they will have to look at samples_pulled counter instead.

What do you think? Would that make the logic clear and much simpler to work with?

I think it'd also make it much easier to look up data in the dataset if we want to see the tokens - because the sample indices will be absolute.

[jaketae]

Yes, this totally makes sense to me. IMHO keeping everything relative is the source of confusion and potential bugs. Of course, there is also an advantage of using relative intervals instead of absolute sample indices, but the more general use case seems to favor the latter. I'm assuming we know what the specific index of the data is that caused a spike.

One question is, when you say index, is this in the sense of Dataset.__getitem__ or the batch index (not sure if this is rigorously defined in ML lingo, but basically the ith batch that is returned after i iter calls to the data loader). I assume the batch gets sampled and shuffled; does this present any complication?

[stas00]

The problem is that the current implementation may potentially lead to the same samples being fed through multiple times in a single epoch.

I was thinking of sample id being the index, because this is what the user sees.

Remember we have batch size rampup, so if we don't use an absolute sample counter we again have a bit of an issue of calculating the sample id.

[jaketae]

Ah right, I missed the batch size ramp up part. I'll try to open a PR next week and maybe we can go from there? I assume this isn't needed as of immediate, but it might be good to have once the training starts and we run into an unrecoverable instability (fingers crossed we won't).

[stas00]

Well, the current solution is sort of working at the point of wanting to skip - the main concern is potentially multiple feedings of the same data. So it's hard to tell what way to proceed, since I'd rather not test something new when the training has started.

I wonder if we should just leave it as it is for now.