Add loss generating token counts #1610
Conversation
|
@mvpatel2000 yeah, CI won't pass until release |
| @@ -1310,9 +1321,11 @@ def build_from_hf( | |||
| raise NotImplementedError() | |||
|
|
|||
| batch_collated = dl.dataloader.collate_fn(batch_tokenized) # type: ignore | |||
| actual_token_count = dl.get_num_tokens_in_batch(batch_collated) | |||
| actual_total_token_count = dl.get_num_tokens_in_batch(batch_collated, token_type='total') | |||
There was a problem hiding this comment.
i might be missing something, but how can we pass in token_type here when it's not in the function definition here? https://github.com/mosaicml/llm-foundry/pull/1610/files#diff-9568d89aed75ca69416abe2a592c6bb9732129049a62c34e4e9263c18495a236R99
There was a problem hiding this comment.
The function being called here is actually defined on the DataSpec class in Composer (https://github.com/mosaicml/composer/blob/28756dd52e96371689b764cb72c336406460ad35/composer/core/data_spec.py#L301). The DataSpec takes in a function from the user and uses it.
There was a problem hiding this comment.
Part of the reason for doing it this way was to maintain backwards compatibility with any existing user defined get_num_tokens_in_batch functions out there.
| torch.sum(batch['labels'] != CROSS_ENTROPY_IGNORE_INDEX).item(), | ||
| ) | ||
|
|
||
| # Subtract one for each example in the batch that starts with a non -100, |
There was a problem hiding this comment.
@dakinggg I don't think this subtraction isn't necessary. Instead you can just do this:
loss_generating_tokens = int(
torch.sum(batch['labels'][...,1:] != CROSS_ENTROPY_IGNORE_INDEX).item(),
)*I just came across this pr while looking into how mosaic's libs handle the gradient accumulation bug recently discussed on x.com
There was a problem hiding this comment.
ah yeah, that should work too :)
Takes advantage of the new functionality in Composer to weight microbatches by loss generating tokens, and not just total tokens. See the Composer PR (mosaicml/composer#3677) for more details and manual testing.
Note: this needs a Composer release and bump (and CI won't pass until that happens)