reduce GPU memory by keeping necessary modules of query_enc

save memory usage

some quick fix

code clean up

some quick fix

trim checkpoint model weights

implement gpt-k demo