[peculator training] Update benchmark_speculator_logical.py to support gpt_bigcode/granite

[sahilsuneja1]

@daviswer: should we also add the caller script?

[daviswer]

Thanks @sahilsuneja1 for putting this together! I don't think we need the caller script - it's ultimately just a simple python call with a bunch of arguments, right? I'd just add a comment with a sample call or two (Llama 7b / granite 20b?) to the top of the benchmark_speculator_logical.py file

[sahilsuneja1]

Updated formatting

[daviswer]

Further CI complaints. I'll try and figure out how to get them running automatically for you rather than having to wait for my explicit go-ahead

[sahilsuneja1]

Fixed isort issues

[sahilsuneja1]

Dunno how to fix the mypy errors

[nairbv]

There are some suggestions here on how to fix some of these, possibly via a fully qualified import? https://stackoverflow.com/questions/68695851/mypy-cannot-find-implementation-or-library-stub-for-module

@afrittoli may also have suggestions.

pasting the mypy errors here, for easier reference:

speculator/benchmark_speculator_logical.py:7: error: Cannot find implementation or library stub for module named "fms_extras.models.paged_gpt_bigcode"  [import-not-found]
speculator/benchmark_speculator_logical.py:7: error: Cannot find implementation or library stub for module named "fms_extras.models"  [import-not-found]
speculator/benchmark_speculator_logical.py:7: error: Cannot find implementation or library stub for module named "fms_extras"  [import-not-found]
speculator/benchmark_speculator_logical.py:8: error: Cannot find implementation or library stub for module named "fms_extras.models.paged_llama"  [import-not-found]
speculator/benchmark_speculator_logical.py:13: error: Cannot find implementation or library stub for module named "fms_extras.models.speculator"  [import-not-found]
speculator/benchmark_speculator_logical.py:14: error: Cannot find implementation or library stub for module named "fms_extras.utils.generation"  [import-not-found]
speculator/benchmark_speculator_logical.py:16: error: Library stubs not installed for "tqdm"  [import-untyped]
speculator/benchmark_speculator_logical.py:173: error: Incompatible types in assignment (expression has type "None", variable has type "str")  [assignment]
speculator/benchmark_speculator_logical.py:208: error: Cannot find implementation or library stub for module named "fms_extras.utils.cache.paged"  [import-not-found]
speculator/benchmark_speculator_logical.py:243: error: Need type annotation for "data" (hint: "data: List[<type>] = ...")  [var-annotated]
speculator/benchmark_speculator_logical.py:247: error: No overload variant of "next" matches argument type "Streaming_Doc_Dataset"  [call-overload]
speculator/benchmark_speculator_logical.py:247: note: Possible overload variants:
speculator/benchmark_speculator_logical.py:247: note:     def [_T] next(SupportsNext[_T], /) -> _T
speculator/benchmark_speculator_logical.py:247: note:     def [_T, _VT] next(SupportsNext[_T], _VT, /) -> _T | _VT
speculator/benchmark_speculator_logical.py:254: error: Incompatible types in assignment (expression has type "Tensor", variable has type "list[Any]")  [assignment]
speculator/benchmark_speculator_logical.py:340: error: Incompatible types in assignment (expression has type "Any | None", variable has type "Callable[..., Any]")  [assignment]
speculator/benchmark_speculator_logical.py:344: error: "list[Any]" has no attribute "size"  [attr-defined]
speculator/benchmark_speculator_logical.py:347: error: No overload variant of "__getitem__" of "list" matches argument type "tuple[slice, slice]"  [call-overload]
speculator/benchmark_speculator_logical.py:347: note: Possible overload variants:
speculator/benchmark_speculator_logical.py:347: note:     def __getitem__(self, SupportsIndex, /) -> Any
speculator/benchmark_speculator_logical.py:347: note:     def __getitem__(self, slice, /) -> list[Any]
speculator/benchmark_speculator.py:6: error: Cannot find implementation or library stub for module named "fms_extras.models.paged_llama"  [import-not-found]
speculator/benchmark_speculator.py:6: error: Cannot find implementation or library stub for module named "fms_extras.models"  [import-not-found]
speculator/benchmark_speculator.py:6: error: Cannot find implementation or library stub for module named "fms_extras"  [import-not-found]
speculator/benchmark_speculator.py:11: error: Cannot find implementation or library stub for module named "fms_extras.models.speculator"  [import-not-found]
speculator/benchmark_speculator.py:12: error: Cannot find implementation or library stub for module named "fms_extras.utils.generation"  [import-not-found]
speculator/benchmark_speculator.py:14: error: Library stubs not installed for "tqdm"  [import-untyped]
speculator/benchmark_speculator.py:14: note: Hint: "python3 -m pip install types-tqdm"
speculator/benchmark_speculator.py:14: note: (or run "mypy --install-types" to install all missing stub packages)
speculator/benchmark_speculator.py:[132](https://github.com/foundation-model-stack/fms-fsdp/actions/runs/8456415698/job/23168426272?pr=62#step:6:133): error: Incompatible types in assignment (expression has type "None", variable has type "str")  [assignment]
speculator/benchmark_speculator.py:164: error: Cannot find implementation or library stub for module named "fms_extras.utils.cache.paged"  [import-not-found]
speculator/benchmark_speculator.py:164: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
speculator/benchmark_speculator.py:193: error: Need type annotation for "data" (hint: "data: List[<type>] = ...")  [var-annotated]
speculator/benchmark_speculator.py:197: error: No overload variant of "next" matches argument type "Streaming_Doc_Dataset"  [call-overload]
speculator/benchmark_speculator.py:197: note: Possible overload variants:
speculator/benchmark_speculator.py:197: note:     def [_T] next(SupportsNext[_T], /) -> _T
speculator/benchmark_speculator.py:197: note:     def [_T, _VT] next(SupportsNext[_T], _VT, /) -> _T | _VT
speculator/benchmark_speculator.py:204: error: Incompatible types in assignment (expression has type "Tensor", variable has type "list[Any]")  [assignment]
speculator/benchmark_speculator.py:281: error: Incompatible types in assignment (expression has type "Any | None", variable has type "Callable[..., Any]")  [assignment]
speculator/benchmark_speculator.py:285: error: "list[Any]" has no attribute "size"  [attr-defined]
speculator/train_speculator.py:4: error: Skipping analyzing "fire": module is installed, but missing library stubs or py.typed marker  [import-untyped]
speculator/train_speculator.py:9: error: Cannot find implementation or library stub for module named "fms_extras.models.speculator"  [import-not-found]
speculator/train_speculator.py:14: error: Cannot find implementation or library stub for module named "train_speculator_utils"  [import-not-found]
Found 31 errors in 3 files (checked 18 source files)

[daviswer]

Looks like a bunch of that is because this is relying on the paged attention branch, which hasn't fully landed in fms-extras/main yet

[daviswer]

Actually it looks like we don't have any dependency on fms-extras listed in requirements - I'm now hitting the same issue with #35. If we can add that then this should go through. @lchu-ibm or @nairbv , do you know the right way to get fms-extras into requirements.txt?

[nairbv]

Actually it looks like we don't have any dependency on fms-extras listed in requirements - I'm now hitting the same issue with #35. If we can add that then this should go through. @lchu-ibm or @nairbv , do you know the right way to get fms-extras into requirements.txt?

we can either point directly at the github URL within requirements.txt (see some examples here: https://stackoverflow.com/questions/16584552/how-to-state-in-requirements-txt-a-direct-github-source) or we could publish fms-extras to pypi. probably the first option is simpler/easier