Fix mirostat state when using multiple sequences

[KerfuffleV2]

As mentioned in #3537, mirostat currently isn't compatible with using multiple sequences.

The main selling point of the way this pull is implemented is such a way that it's pretty simple and uninvasive.

However, I really don't like storing mutable sampler state in gpt_params (even though it's only in common and not the main llama.cpp API). This also required removing the const from the params argument to llama_sample_token. As far as I can see, the existing examples don't care about that.

I feel like the right way to do this is probably to move the sampler state out of gpt_params and have it passed separately. In that case, this is probably also where grammar should be since it's a type of sampler state. So we wouldn't add a new argument to llama_sample_token, we'd replace the current grammar one with sampler state. This of course would require changing a lot more stuff, including all the examples that use llama_sample_token (I don't think it would be too bad though).

Thoughts?

Closes #3537

[ggerganov]

Yes, dedicated sampling state with grammar and mirostat would be better. Maybe implemented in common/sampling.h/.cpp. It should probably inherit all sampling-related parameters gpt_params, such as temperature, top_p, top_k, etc, so that llama_sample_token accepts struct llama_sampling_state instead of struct gpt_params.

For now we can have this workaround

[KerfuffleV2]

@ggerganov

For now we can have this workaround

Do you actually prefer doing it this way for now?

I don't mind changing this to do it the other way I suggested as long as you agree that approach is okay.

Maybe implemented in common/sampling.h/.cpp.

This actually raises another question that I'm actually dealing with in my seqrep sampler. Right now it's really awkward to have multiple source files in common. This is how I dealt with it:

COMMON_DEPS = common/common.cpp common/common.h build-info.h common/log.h
COMMON_OBJS = common.o
ifndef LLAMA_DISABLE_SEQREP_SAMPLER
COMMON_DEPS += common/seqrep-sampler.cpp common/seqrep-sampler.h
COMMON_OBJS += seqrep-sampler.o
endif
common.o: $(COMMON_DEPS)
	$(CXX) $(CXXFLAGS) -c $< -o $@

simple: examples/simple/simple.cpp                            build-info.h ggml.o llama.o $(COMMON_OBJS) $(OBJS)
	$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)

and so on. It's probably less of a pain in cmake.

[ggerganov]

We should try separating the sampling state - would be better than the current fix, so let's give it a try if you are up to.

The proposed Makefile looks OK to me.

[FSSRepo]

When this commit is merged into the master branch, we need to add the parameter slot.id in llama_sample_token(ctx, NULL, NULL, params, slot.tokens_prev, candidates, slot.i_batch, slot.id) in the llama_sample_token call of the server-parallel in #3490 and this params.sampler_state.erase(slot.id); too.

[KerfuffleV2]

Well, this went from small and self-contained to huge and complicated. I sure hope I'm on the right track after all this.

Pretty much all the sampling stuff got moved into common/sampling.{c,h}. Creating sampling state takes gpt_params and llama_grammar * (can be NULL). The sampling-related params in gpt_params are also in a llama_sampling_params struct now.

The sampling state (llama_sampling_state) holds a copy of the params that were created at init time. It also holds per-sequence state for samplers that use it (Mirostat 1/2, and grammar).

If the per-sequence state doesn't exist when llama_sample_token is called, it will be created with default values. In the case of grammar, this does llama_grammar_copy on the llama_grammar * that was supplied at init time.

If you want to use separate grammar, or separate grammar states per sequence then you'd probably have to manually manage the grammar part yourself. (Not sure this is so easy right now, so I might need to add an interface.)

In the case of parallel generation, when a sequence ends and you will want to reuse it (or just free up memory) you need to call llama_sampling_state_reset with the sequence id. This will reset stuff like Mirostat mu to the initial value so it won't mess with future generation with that sequence id.

I also randomly threw in support for LLAMA_SANITIZE_{THREAD,ADDRESS,UNDEFINED} in the Makefile.

I think this currently doesn't break stuff, but there were some tricky parts like server and speculative.

[KerfuffleV2]

@FSSRepo You're going to hate me when you see the next step in this pull.

Calling sampling is going to look like:

llama_sample_token(ctx, NULL, sampling_state, slot.tokens_prev, candidates, slot.i_batch, slot.id);

It looks like #3490 doesn't support grammar currently? So that's going to make your life easier. Pretty much the only other thing to worry about is calling llama_sampling_state_reset when you're done with a sequence id (but may want to reuse it for a different generation). Pretty much when you hit the EOS token or reach the quota of tokens to generate, you can just reset that sequence id.

[KerfuffleV2]

I exported the function to fetch/create default instances of sampler state. This should fix the problem I mentioned earlier about how it would be hard to do something like parallel generation where each sequence used its own grammar.

By the way, since the ggml-alloc stuff:

ggml-alloc.c:212:32: runtime error: pointer index expression with base 0x00000100a020 overflowed to 0xffffffffffffffff
llama_new_context_with_model: compute buffer total size = 552.88 MB
llama_new_context_with_model: VRAM scratch buffer: 546.75 MB
llama_new_context_with_model: total VRAM used: 6515.44 MB (model: 3368.69 MB, context: 3146.75 MB)
ggml-cuda.cu:6787:51: runtime error: applying non-zero offset 1152 to null pointer

Not sure if that's anything to worry about. The ggml-alloc one is new, the other one isn't. (I was assuming it was just because GCC's address sanitizing stuff doesn't know about CUDA/ROCM.) edit: Might not be a problem since it seems like it's triggering just on doing pointer math that results in something out of bounds rather than actually accessing it.

[KerfuffleV2]

This isn't really approved, right? Even without a full review, any changes I can/should start working on?

[ggerganov]

Should we absorb last_tokens and candidates into llama_sampling_state?

[KerfuffleV2]

Last tokens can be specific to the sequence, right? So this would kind of mean stuff using last tokens would have to be aware of sequences. I kind of feel like this also might limit how people can manipulate last tokens and if there are currently examples that do that kind of thing it might be difficult to adapt them (for me anyway, since I'm not really deeply familiar with most of them).

candidates I'm less sure about, it's basically just a scratch area for the logits in a form samplers can work with (right?) so I think moving it in there is less of a big deal. It's a pretty large structure though, I don't know if that's a consideration. Right now the current stuff in those structs is pretty lightweight.

I don't have very strong feelings about this. I'd like to say "These changes are already complicated enough, let's come back to that" but... I probably never would. :)

[ggerganov]

Probably we should add llama_sampling_accept_token() and move this call in there together with update of last_tokens member (if we decide it should become part of llama_sampling_state)

[KerfuffleV2]

This might be a bit of a pain since it depends on a number of other static functions like decode_utf8, llama_grammar_accept, llama_token_to_str. It looks like the grammar stuff is the only thing that uses them, so maybe they could be moved too. I'm not sure what other parts of the grammar code depends on them though, so it might not be that simple.

[ggerganov]

This is a great change. I actually think that we should merge llama_sampling directly into llama.cpp. But let's do this after this PR is merged and tested for some time

Not sure if that's anything to worry about.

These errors look benign, but we will look in ways to fix them anyway

[KerfuffleV2]

Current status: I took the suggestions to rename the functions/types but I didn't do stuff like moving last_tokens into the sampling context (yet). edit: Just to be clear, the "(yet)" doesn't mean I'm actually planning to unless someone insists on it.

[ggerganov]

Suggested change

	struct llama_sampling_context & sampling_ctx,
	struct llama_sampling_context & ctx_sampling,

There are a few other places that need similar change for consistency sake

[KerfuffleV2]

By "a few" you mean 20 or so? :) Hopefully I caught them all. Everything seems to compile/work still.

[ggerganov]

Thanks for this - I accidentally merged this too quickly with the old title. Should have updated to the more relevant change of introducing the llama_sampling_context.

[ggerganov]

@KerfuffleV2

Any reason not to have single sequence data in llama_sampling_context?
When we want to sample multiple sequences, we will create on llama_sampling_context for each.

This way, each sequence can also have a separate llama_grammar instance which seems to make sense.

I.e. have it like this:

// general sampler context
typedef struct llama_sampling_context {
    ~llama_sampling_context();

    // parameters that will be used for sampling
    llama_sampling_params params;

    float mirostat_mu; // mirostat sampler state

    llama_grammar * grammar;
} llama_sampling_context;