search

SulRash / envenc

Repository for environment encoder, an attempt at improving reinforcement learning agents' generalisability through learning how to act on universal multimodal embeddings generated by a vision-language model.
MIT License
2 stars 0 forks source link

Update Repository to Torch 2.0 #5

Closed SulRash closed 4 months ago

SulRash commented 4 months ago

Updating to torch 2 would be so nice. I'd be able to use torch.compile and other features to improve throughput, the VLM is such a severe bottleneck...

SulRash commented 4 months ago

Just to remember, steps per second (SPS) used to be 13

SulRash commented 4 months ago

With torch.compile and torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False) I was able to get the SPS to ~25

SulRash commented 4 months ago

I also loaded in the model in 4bit quantisations with bits and bytes, but the model is so small I'm worried this will degrade performance massively...

SulRash commented 4 months ago

For some reason it keeps failing when I try to use bitsandbytes + torch.compile

SulRash commented 4 months ago

I removed quantisation for now

SulRash commented 4 months ago

I should switch to using model.generate while adding output_hidden_states as a kwarg instead of just going through the forward, might help with doing batched inference over multiple environments.

SulRash commented 4 months ago

This scope is getting too big, update is complete ill make another branch for improving llm performance with batching and stuff