>>102188727
Try using a Nemo quant, I managed to make it work with this change:
# Create the draft model class
class DraftModel(LlamaDraftModel):
def __init__(self, current_model, path_to_draft_model, n_speculation_tokens):
self.model = Llama(
verbose=True,
model_path=args.model_draft,
n_gpu_layers=args.n_gpu_layers_draft,
n_ctx=args.ctx_size,
n_batch=args.batch_size,
n_threads=args.threads_draft,
n_threads_batch=args.threads_batch_draft,
flash_attn=args.flash_attn,
numa=numa_strategy,
use_mmap=not args.no_mmap,
use_mlock=args.mlock
)
self.n_speculation_tokens = n_speculation_tokens
self.current_model = current_model
def __call__(self, input_ids, **kwargs):
text = self.current_model.detokenize(input_ids)
generator = self.model.generate(self.model.tokenize(text), top_k=1, temp=0, top_p=1)
output = np.zeros(self.n_speculation_tokens, dtype=np.intc)
for i in range(self.n_speculation_tokens):
output[i] = self.current_model.tokenize(self.model.detokenize([next(generator)]))[1]
return output
# Set the custom draft model
llama_proxy = next(get_llama_proxy())
llama_model = llama_proxy._current_model
if llama_model is not None:
llama_model.draft_model = DraftModel(llama_model, args.model_draft, args.draft)
For me the Nemo draft model seems to get the right tokens relatively often.
