add support for H100s

[thelinuxkid]

What does this PR do?

add support for H100s

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Was this discussed/approved via a Github issue or the discord / slack channel? Please add a link
  to it if that's the case.
• Did you write any new necessary tests?

Who can review?

@arnavgarg1
@tgaddair

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

POC

[flozi00]

https://github.com/enricai/lorax/blob/58b3c0c7a97f3e43355af4965afc09668e58f571/Dockerfile#L40

Could you set the pytorch version in docker to 2.1.1 too ?

[thelinuxkid]

https://github.com/enricai/lorax/blob/58b3c0c7a97f3e43355af4965afc09668e58f571/Dockerfile#L40

Could you set the pytorch version in docker to 2.1.1 too ?

@flozi00 Done!!

[thelinuxkid]

Build issue should be fixed now

[tgaddair]

LGTM! I'll go ahead and land, and then do some robustness tests after the new image goes out.

[tgaddair]

@thelinuxkid looks like based on the tests we actually need to add back peft as a dependency.

[tgaddair]

@thelinuxkid alternatively, would you mind giving me edit access to the branch / PR so I can make these changes while I test things out?

[thelinuxkid]

@tgaddair just added, will wait for tests to see if they turn out OK this time

[tgaddair]

Hey @thelinuxkid, apologies, but can you add these to pyproject.toml as well and update the other dep files:

boto3 = "^1.28.34"
urllib3 = "<=1.26.18"

[thelinuxkid]

@tgaddair Done! Also gave you access as maintainer to our fork

[tgaddair]

Ran a few tests earlier today with these changes. Everything looks good on my end. Merging!

[thelinuxkid]

Thanks @tgaddair ! When will this be available in the docker image? I'm not able to run with flash attn successfully. I build it but I still get an error that it doesn't exist when running it