bigscience-workshop / bigscience-workshop/Megatron-DeepSpeed

Calling IndexedDatasetBuilder directly with a best_fit datatype fails

Open
#67 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.4k
Forks
226
PR merge metrics
No merged PRs in 30d

Description

I don't know whether this is intended to work or not, but I found the following program:
```
from megatron.data.indexed_dataset import IndexedDatasetBuilder, best_fitting_dtype

best_dtype = best_fitting_dtype(10_000)
IndexedDatasetBuilder("testfile", dtype=best_dtype)
```
leads to an error like:
```
File "/path/to/Megatron-DeepSpeed.git/megatron/data/indexed_dataset.py", line 284, in __init__
self.element_size = self.element_sizes[self.dtype]
KeyError:
```
This shows up because ``best_fitting_dtype`` will return ``numpy.uint16`` for small vocabs:

https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/c680954e0b13232abd5b72711f2032bea1ad65c9/megatron/data/indexed_dataset.py#L25-L27

but that particular type is missing from the ``element_sizes`` table.

https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/c680954e0b13232abd5b72711f2032bea1ad65c9/megatron/data/indexed_dataset.py#L268-L276

which is used in the ``IndexedDatasetBuilder`` constructor here:

https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/c680954e0b13232abd5b72711f2032bea1ad65c9/megatron/data/indexed_dataset.py#L284

Should something like this work?

If so, it seems like either the ``best_fitting_dtype`` function should return a different type or ``uint16`` should be added to the ``element_sizes`` table like https://github.com/bigscience-workshop/Megatron-DeepSpeed/pull/55/commits/f706108acadb2d97dc77f066930accfbdde4942f.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.